refactor h2 dispatcher to async/await.reduce duplicate code

actix-http/src/h2/dispatcher.rs

@@ -3,18 +3,16 @@ use std::{cmp, future::Future, marker::PhantomData, net, pin::Pin, rc::Rc};
 
 use actix_codec::{AsyncRead, AsyncWrite};
 use actix_service::Service;
+use actix_utils::future::poll_fn;
 use bytes::{Bytes, BytesMut};
 use futures_core::ready;
-use h2::{
+use h2::server::{Connection, SendResponse};
-    server::{Connection, SendResponse},
-    SendStream,
-};
 use http::header::{HeaderValue, CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING};
 use log::{error, trace};
 
-use crate::body::{Body, BodySize, MessageBody};
+use crate::body::{BodySize, MessageBody};
 use crate::config::ServiceConfig;
-use crate::error::{DispatchError, Error};
+use crate::error::Error;
 use crate::message::ResponseHead;
 use crate::payload::Payload;
 use crate::request::Request;
@@ -71,14 +69,14 @@ where
     T: AsyncRead + AsyncWrite + Unpin,
 
     S: Service<Request>,
-    S::Error: Into<Error> + 'static,
+    S::Error: Into<Error>,
     S::Future: 'static,
-    S::Response: Into<Response<B>> + 'static,
+    S::Response: Into<Response<B>>,
 
-    B: MessageBody + 'static,
+    B: MessageBody,
     B::Error: Into<Error>,
 {
-    type Output = Result<(), DispatchError>;
+    type Output = Result<(), crate::error::DispatchError>;
 
     #[inline]
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
@@ -90,7 +88,7 @@ where
 
                 Some(Err(err)) => return Poll::Ready(Err(err.into())),
 
-                Some(Ok((req, res))) => {
+                Some(Ok((req, tx))) => {
                     let (parts, body) = req.into_parts();
                     let pl = crate::h2::Payload::new(body);
                     let pl = Payload::<crate::payload::PayloadStream>::H2(pl);
@@ -106,296 +104,178 @@ where
                     // merge on_connect_ext data into request extensions
                     this.on_connect_data.merge_into(&mut req);
 
-                    let svc = ServiceResponse {
+                    let fut = this.flow.service.call(req);
-                        state: ServiceResponseState::ServiceCall(
+                    let config = this.config.clone();
-                            this.flow.service.call(req),
-                            Some(res),
-                        ),
-                        config: this.config.clone(),
-                        buffer: None,
-                        _phantom: PhantomData,
-                    };
 
-                    actix_rt::spawn(svc);
+                    // multiplex request handling with spawn task
+                    actix_rt::spawn(async move {
+                        // resolve service call and send response.
+                        let res = match fut.await {
+                            Ok(res) => {
+                                let (res, body) = res.into().replace_body(());
+                                handle_response(res, body, tx, config).await
+                            }
+                            Err(err) => {
+                                let (res, body) =
+                                    Response::from_error(err.into()).replace_body(());
+                                handle_response(res, body, tx, config).await
+                            }
+                        };
+
+                        // log error.
+                        if let Err(err) = res {
+                            match err {
+                                DispatchError::SendResponse(err) => {
+                                    trace!("Error sending HTTP/2 response: {:?}", err)
+                                }
+                                DispatchError::SendData(err) => warn!("{:?}", err),
+                                DispatchError::ResponseBody(err) => {
+                                    error!("Response payload stream error: {:?}", err)
+                                }
+                            }
+                        }
+                    });
                 }
             }
         }
     }
 }
 
-#[pin_project::pin_project]
+enum DispatchError {
-struct ServiceResponse<F, I, E, B> {
+    SendResponse(h2::Error),
-    #[pin]
+    SendData(h2::Error),
-    state: ServiceResponseState<F, B>,
+    ResponseBody(Error),
+}
+
+async fn handle_response<B>(
+    res: Response<()>,
+    body: B,
+    mut tx: SendResponse<Bytes>,
     config: ServiceConfig,
-    buffer: Option<Bytes>,
+) -> Result<(), DispatchError>
-    _phantom: PhantomData<(I, E)>,
-}
-
-#[pin_project::pin_project(project = ServiceResponseStateProj)]
-enum ServiceResponseState<F, B> {
-    ServiceCall(#[pin] F, Option<SendResponse<Bytes>>),
-    SendPayload(SendStream<Bytes>, #[pin] B),
-    SendErrorPayload(SendStream<Bytes>, #[pin] Body),
-}
-
-impl<F, I, E, B> ServiceResponse<F, I, E, B>
 where
-    F: Future<Output = Result<I, E>>,
-    E: Into<Error>,
-    I: Into<Response<B>>,
-
     B: MessageBody,
     B::Error: Into<Error>,
 {
-    fn prepare_response(
+    // prepare response.
-        &self,
+    let mut size = body.size();
-        head: &ResponseHead,
+    let res = prepare_response(config, res.head(), &mut size);
-        size: &mut BodySize,
+    let eof = size.is_eof();
-    ) -> http::Response<()> {
-        let mut has_date = false;
-        let mut skip_len = size != &BodySize::Stream;
 
-        let mut res = http::Response::new(());
+    // send response head and return on eof.
-        *res.status_mut() = head.status;
+    let mut stream = tx
-        *res.version_mut() = http::Version::HTTP_2;
+        .send_response(res, eof)
+        .map_err(DispatchError::SendResponse)?;
 
-        // Content length
+    if eof {
-        match head.status {
+        return Ok(());
-            http::StatusCode::NO_CONTENT
+    }
-            | http::StatusCode::CONTINUE
+
-            | http::StatusCode::PROCESSING => *size = BodySize::None,
+    // poll response body and send chunks to client.
-            http::StatusCode::SWITCHING_PROTOCOLS => {
+    actix_rt::pin!(body);
-                skip_len = true;
+
-                *size = BodySize::Stream;
+    while let Some(res) = poll_fn(|cx| body.as_mut().poll_next(cx)).await {
+        let mut chunk = res.map_err(|err| DispatchError::ResponseBody(err.into()))?;
+
+        'send: loop {
+            // reserve enough space and wait for stream ready.
+            stream.reserve_capacity(cmp::min(chunk.len(), CHUNK_SIZE));
+
+            match poll_fn(|cx| stream.poll_capacity(cx)).await {
+                // No capacity left. drop body and return.
+                None => return Ok(()),
+                Some(res) => {
+                    // Split chuck to writeable size and send to client.
+                    let cap = res.map_err(DispatchError::SendData)?;
+
+                    let len = chunk.len();
+                    let bytes = chunk.split_to(cmp::min(cap, len));
+
+                    stream
+                        .send_data(bytes, false)
+                        .map_err(DispatchError::SendData)?;
+
+                    // Current chuck completely sent. break send loop and poll next one.
+                    if chunk.is_empty() {
+                        break 'send;
+                    }
+                }
             }
+        }
+    }
+
+    // response body streaming finished. send end of stream and return.
+    stream
+        .send_data(Bytes::new(), true)
+        .map_err(DispatchError::SendData)?;
+
+    Ok(())
+}
+
+fn prepare_response(
+    config: ServiceConfig,
+    head: &ResponseHead,
+    size: &mut BodySize,
+) -> http::Response<()> {
+    let mut has_date = false;
+    let mut skip_len = size != &BodySize::Stream;
+
+    let mut res = http::Response::new(());
+    *res.status_mut() = head.status;
+    *res.version_mut() = http::Version::HTTP_2;
+
+    // Content length
+    match head.status {
+        http::StatusCode::NO_CONTENT
+        | http::StatusCode::CONTINUE
+        | http::StatusCode::PROCESSING => *size = BodySize::None,
+        http::StatusCode::SWITCHING_PROTOCOLS => {
+            skip_len = true;
+            *size = BodySize::Stream;
+        }
+        _ => {}
+    }
+
+    let _ = match size {
+        BodySize::None | BodySize::Stream => None,
+        BodySize::Empty => res
+            .headers_mut()
+            .insert(CONTENT_LENGTH, HeaderValue::from_static("0")),
+        BodySize::Sized(len) => {
+            let mut buf = itoa::Buffer::new();
+
+            res.headers_mut().insert(
+                CONTENT_LENGTH,
+                HeaderValue::from_str(buf.format(*len)).unwrap(),
+            )
+        }
+    };
+
+    // copy headers
+    for (key, value) in head.headers.iter() {
+        match *key {
+            // TODO: consider skipping other headers according to:
+            //       https://tools.ietf.org/html/rfc7540#section-8.1.2.2
+            // omit HTTP/1.x only headers
+            CONNECTION | TRANSFER_ENCODING => continue,
+            CONTENT_LENGTH if skip_len => continue,
+            DATE => has_date = true,
             _ => {}
         }
 
-        let _ = match size {
+        res.headers_mut().append(key, value.clone());
-            BodySize::None | BodySize::Stream => None,
-            BodySize::Empty => res
-                .headers_mut()
-                .insert(CONTENT_LENGTH, HeaderValue::from_static("0")),
-            BodySize::Sized(len) => {
-                let mut buf = itoa::Buffer::new();
-
-                res.headers_mut().insert(
-                    CONTENT_LENGTH,
-                    HeaderValue::from_str(buf.format(*len)).unwrap(),
-                )
-            }
-        };
-
-        // copy headers
-        for (key, value) in head.headers.iter() {
-            match *key {
-                // TODO: consider skipping other headers according to:
-                //       https://tools.ietf.org/html/rfc7540#section-8.1.2.2
-                // omit HTTP/1.x only headers
-                CONNECTION | TRANSFER_ENCODING => continue,
-                CONTENT_LENGTH if skip_len => continue,
-                DATE => has_date = true,
-                _ => {}
-            }
-
-            res.headers_mut().append(key, value.clone());
-        }
-
-        // set date header
-        if !has_date {
-            let mut bytes = BytesMut::with_capacity(29);
-            self.config.set_date_header(&mut bytes);
-            res.headers_mut().insert(
-                DATE,
-                // SAFETY: serialized date-times are known ASCII strings
-                unsafe { HeaderValue::from_maybe_shared_unchecked(bytes.freeze()) },
-            );
-        }
-
-        res
     }
-}
 
-impl<F, I, E, B> Future for ServiceResponse<F, I, E, B>
+    // set date header
-where
+    if !has_date {
-    F: Future<Output = Result<I, E>>,
+        let mut bytes = BytesMut::with_capacity(29);
-    E: Into<Error>,
+        config.set_date_header(&mut bytes);
-    I: Into<Response<B>>,
+        res.headers_mut().insert(
-
+            DATE,
-    B: MessageBody,
+            // SAFETY: serialized date-times are known ASCII strings
-    B::Error: Into<Error>,
+            unsafe { HeaderValue::from_maybe_shared_unchecked(bytes.freeze()) },
-{
+        );
-    type Output = ();
-
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let mut this = self.as_mut().project();
-
-        match this.state.project() {
-            ServiceResponseStateProj::ServiceCall(call, send) => {
-                match ready!(call.poll(cx)) {
-                    Ok(res) => {
-                        let (res, body) = res.into().replace_body(());
-
-                        let mut send = send.take().unwrap();
-                        let mut size = body.size();
-                        let h2_res =
-                            self.as_mut().prepare_response(res.head(), &mut size);
-                        this = self.as_mut().project();
-
-                        let stream = match send.send_response(h2_res, size.is_eof()) {
-                            Err(e) => {
-                                trace!("Error sending HTTP/2 response: {:?}", e);
-                                return Poll::Ready(());
-                            }
-                            Ok(stream) => stream,
-                        };
-
-                        if size.is_eof() {
-                            Poll::Ready(())
-                        } else {
-                            this.state
-                                .set(ServiceResponseState::SendPayload(stream, body));
-                            self.poll(cx)
-                        }
-                    }
-
-                    Err(err) => {
-                        let res = Response::from_error(err.into());
-                        let (res, body) = res.replace_body(());
-
-                        let mut send = send.take().unwrap();
-                        let mut size = body.size();
-                        let h2_res =
-                            self.as_mut().prepare_response(res.head(), &mut size);
-                        this = self.as_mut().project();
-
-                        let stream = match send.send_response(h2_res, size.is_eof()) {
-                            Err(e) => {
-                                trace!("Error sending HTTP/2 response: {:?}", e);
-                                return Poll::Ready(());
-                            }
-                            Ok(stream) => stream,
-                        };
-
-                        if size.is_eof() {
-                            Poll::Ready(())
-                        } else {
-                            this.state.set(ServiceResponseState::SendErrorPayload(
-                                stream, body,
-                            ));
-                            self.poll(cx)
-                        }
-                    }
-                }
-            }
-
-            ServiceResponseStateProj::SendPayload(ref mut stream, ref mut body) => {
-                loop {
-                    match this.buffer {
-                        Some(ref mut buffer) => match ready!(stream.poll_capacity(cx)) {
-                            None => return Poll::Ready(()),
-
-                            Some(Ok(cap)) => {
-                                let len = buffer.len();
-                                let bytes = buffer.split_to(cmp::min(cap, len));
-
-                                if let Err(e) = stream.send_data(bytes, false) {
-                                    warn!("{:?}", e);
-                                    return Poll::Ready(());
-                                } else if !buffer.is_empty() {
-                                    let cap = cmp::min(buffer.len(), CHUNK_SIZE);
-                                    stream.reserve_capacity(cap);
-                                } else {
-                                    this.buffer.take();
-                                }
-                            }
-
-                            Some(Err(e)) => {
-                                warn!("{:?}", e);
-                                return Poll::Ready(());
-                            }
-                        },
-
-                        None => match ready!(body.as_mut().poll_next(cx)) {
-                            None => {
-                                if let Err(e) = stream.send_data(Bytes::new(), true) {
-                                    warn!("{:?}", e);
-                                }
-                                return Poll::Ready(());
-                            }
-
-                            Some(Ok(chunk)) => {
-                                stream
-                                    .reserve_capacity(cmp::min(chunk.len(), CHUNK_SIZE));
-                                *this.buffer = Some(chunk);
-                            }
-
-                            Some(Err(err)) => {
-                                error!(
-                                    "Response payload stream error: {:?}",
-                                    err.into()
-                                );
-
-                                return Poll::Ready(());
-                            }
-                        },
-                    }
-                }
-            }
-
-            ServiceResponseStateProj::SendErrorPayload(ref mut stream, ref mut body) => {
-                // TODO: de-dupe impl with SendPayload
-
-                loop {
-                    match this.buffer {
-                        Some(ref mut buffer) => match ready!(stream.poll_capacity(cx)) {
-                            None => return Poll::Ready(()),
-
-                            Some(Ok(cap)) => {
-                                let len = buffer.len();
-                                let bytes = buffer.split_to(cmp::min(cap, len));
-
-                                if let Err(e) = stream.send_data(bytes, false) {
-                                    warn!("{:?}", e);
-                                    return Poll::Ready(());
-                                } else if !buffer.is_empty() {
-                                    let cap = cmp::min(buffer.len(), CHUNK_SIZE);
-                                    stream.reserve_capacity(cap);
-                                } else {
-                                    this.buffer.take();
-                                }
-                            }
-
-                            Some(Err(e)) => {
-                                warn!("{:?}", e);
-                                return Poll::Ready(());
-                            }
-                        },
-
-                        None => match ready!(body.as_mut().poll_next(cx)) {
-                            None => {
-                                if let Err(e) = stream.send_data(Bytes::new(), true) {
-                                    warn!("{:?}", e);
-                                }
-                                return Poll::Ready(());
-                            }
-
-                            Some(Ok(chunk)) => {
-                                stream
-                                    .reserve_capacity(cmp::min(chunk.len(), CHUNK_SIZE));
-                                *this.buffer = Some(chunk);
-                            }
-
-                            Some(Err(err)) => {
-                                error!("Response payload stream error: {:?}", err);
-
-                                return Poll::Ready(());
-                            }
-                        },
-                    }
-                }
-            }
-        }
     }
+
+    res
 }