provide optimisation path for single-chunk body types (#2497)

actix-http/src/body/boxed.rs

@@ -51,6 +51,34 @@ impl MessageBody for BoxBody {
             .poll_next(cx)
             .map_err(|err| Error::new_body().with_cause(err))
     }
+
+    fn is_complete_body(&self) -> bool {
+        self.0.is_complete_body()
+    }
+
+    fn take_complete_body(&mut self) -> Bytes {
+        debug_assert!(
+            self.is_complete_body(),
+            "boxed type does not allow taking complete body; caller should make sure to \
+            call `is_complete_body` first",
+        );
+
+        // we do not have DerefMut access to call take_complete_body directly but since
+        // is_complete_body is true we should expect the entire bytes chunk in one poll_next
+
+        let waker = futures_util::task::noop_waker();
+        let mut cx = Context::from_waker(&waker);
+
+        match self.as_pin_mut().poll_next(&mut cx) {
+            Poll::Ready(Some(Ok(data))) => data,
+            _ => {
+                panic!(
+                    "boxed type indicated it allows taking complete body but failed to \
+                    return Bytes when polled",
+                );
+            }
+        }
+    }
 }
 
 #[cfg(test)]

actix-http/src/body/either.rs

@@ -67,6 +67,20 @@ where
                 .map_err(|err| Error::new_body().with_cause(err)),
         }
     }
+
+    fn is_complete_body(&self) -> bool {
+        match self {
+            EitherBody::Left { body } => body.is_complete_body(),
+            EitherBody::Right { body } => body.is_complete_body(),
+        }
+    }
+
+    fn take_complete_body(&mut self) -> Bytes {
+        match self {
+            EitherBody::Left { body } => body.take_complete_body(),
+            EitherBody::Right { body } => body.take_complete_body(),
+        }
+    }
 }
 
 #[cfg(test)]

actix-http/src/body/message_body.rs

@@ -25,10 +25,58 @@ pub trait MessageBody {
     fn size(&self) -> BodySize;
 
     /// Attempt to pull out the next chunk of body bytes.
+    // TODO: expand documentation
     fn poll_next(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<Bytes, Self::Error>>>;
+
+    /// Returns true if entire body bytes chunk is obtainable in one call to `poll_next`.
+    ///
+    /// This method's implementation should agree with [`take_complete_body`] and should always be
+    /// checked before taking the body.
+    ///
+    /// The default implementation returns `false.
+    ///
+    /// [`take_complete_body`]: MessageBody::take_complete_body
+    fn is_complete_body(&self) -> bool {
+        false
+    }
+
+    /// Returns the complete chunk of body bytes.
+    ///
+    /// Implementors of this method should note the following:
+    /// - It is acceptable to skip the omit checks of [`is_complete_body`]. The responsibility of
+    ///   performing this check is delegated to the caller.
+    /// - If the result of [`is_complete_body`] is conditional, that condition should be given
+    ///   equivalent attention here.
+    /// - A second call call to [`take_complete_body`] should return an empty `Bytes` or panic.
+    /// - A call to [`poll_next`] after calling [`take_complete_body`] should return `None` unless
+    ///   the chunk is guaranteed to be empty.
+    ///
+    /// The default implementation panics unconditionally, indicating a control flow bug in the
+    /// calling code.
+    ///
+    /// # Panics
+    /// With a correct implementation, panics if called without first checking [`is_complete_body`].
+    ///
+    /// [`is_complete_body`]: MessageBody::is_complete_body
+    /// [`take_complete_body`]: MessageBody::take_complete_body
+    /// [`poll_next`]: MessageBody::poll_next
+    fn take_complete_body(&mut self) -> Bytes {
+        assert!(
+            self.is_complete_body(),
+            "type ({}) allows taking complete body but did not provide an implementation \
+            of `take_complete_body`",
+            std::any::type_name::<Self>()
+        );
+
+        unimplemented!(
+            "type ({}) does not allow taking complete body; caller should make sure to \
+            check `is_complete_body` first",
+            std::any::type_name::<Self>()
+        );
+    }
 }
 
 mod foreign_impls {
@@ -49,6 +97,14 @@ mod foreign_impls {
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             match *self {}
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            match *self {}
+        }
     }
 
     impl MessageBody for () {
@@ -66,6 +122,16 @@ mod foreign_impls {
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             Poll::Ready(None)
         }
+
+        #[inline]
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        #[inline]
+        fn take_complete_body(&mut self) -> Bytes {
+            Bytes::new()
+        }
     }
 
     impl<B> MessageBody for Box<B>
@@ -86,6 +152,16 @@ mod foreign_impls {
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             Pin::new(self.get_mut().as_mut()).poll_next(cx)
         }
+
+        #[inline]
+        fn is_complete_body(&self) -> bool {
+            self.as_ref().is_complete_body()
+        }
+
+        #[inline]
+        fn take_complete_body(&mut self) -> Bytes {
+            self.as_mut().take_complete_body()
+        }
     }
 
     impl<B> MessageBody for Pin<Box<B>>
@@ -106,6 +182,38 @@ mod foreign_impls {
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             self.as_mut().poll_next(cx)
         }
+
+        #[inline]
+        fn is_complete_body(&self) -> bool {
+            self.as_ref().is_complete_body()
+        }
+
+        #[inline]
+        fn take_complete_body(&mut self) -> Bytes {
+            debug_assert!(
+                self.is_complete_body(),
+                "inner type \"{}\" does not allow taking complete body; caller should make sure to \
+                call `is_complete_body` first",
+                std::any::type_name::<B>(),
+            );
+
+            // we do not have DerefMut access to call take_complete_body directly but since
+            // is_complete_body is true we should expect the entire bytes chunk in one poll_next
+
+            let waker = futures_util::task::noop_waker();
+            let mut cx = Context::from_waker(&waker);
+
+            match self.as_mut().poll_next(&mut cx) {
+                Poll::Ready(Some(Ok(data))) => data,
+                _ => {
+                    panic!(
+                        "inner type \"{}\" indicated it allows taking complete body but failed to \
+                        return Bytes when polled",
+                        std::any::type_name::<B>()
+                    );
+                }
+            }
+        }
     }
 
     impl MessageBody for &'static [u8] {
@@ -116,17 +224,23 @@ mod foreign_impls {
         }
 
         fn poll_next(
-            self: Pin<&mut Self>,
+            mut self: Pin<&mut Self>,
             _cx: &mut Context<'_>,
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             if self.is_empty() {
                 Poll::Ready(None)
             } else {
-                let bytes = mem::take(self.get_mut());
-                let bytes = Bytes::from_static(bytes);
-                Poll::Ready(Some(Ok(bytes)))
+                Poll::Ready(Some(Ok(self.take_complete_body())))
             }
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            Bytes::from_static(mem::take(self))
+        }
     }
 
     impl MessageBody for Bytes {
@@ -137,16 +251,23 @@ mod foreign_impls {
         }
 
         fn poll_next(
-            self: Pin<&mut Self>,
+            mut self: Pin<&mut Self>,
             _cx: &mut Context<'_>,
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             if self.is_empty() {
                 Poll::Ready(None)
             } else {
-                let bytes = mem::take(self.get_mut());
-                Poll::Ready(Some(Ok(bytes)))
+                Poll::Ready(Some(Ok(self.take_complete_body())))
             }
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            mem::take(self)
+        }
     }
 
     impl MessageBody for BytesMut {
@@ -157,16 +278,23 @@ mod foreign_impls {
         }
 
         fn poll_next(
-            self: Pin<&mut Self>,
+            mut self: Pin<&mut Self>,
             _cx: &mut Context<'_>,
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             if self.is_empty() {
                 Poll::Ready(None)
             } else {
-                let bytes = mem::take(self.get_mut()).freeze();
-                Poll::Ready(Some(Ok(bytes)))
+                Poll::Ready(Some(Ok(self.take_complete_body())))
             }
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            mem::take(self).freeze()
+        }
     }
 
     impl MessageBody for Vec<u8> {
@@ -177,16 +305,23 @@ mod foreign_impls {
         }
 
         fn poll_next(
-            self: Pin<&mut Self>,
+            mut self: Pin<&mut Self>,
             _cx: &mut Context<'_>,
         ) -> Poll<Option<Result<Bytes, Self::Error>>> {
             if self.is_empty() {
                 Poll::Ready(None)
             } else {
-                let bytes = mem::take(self.get_mut());
-                Poll::Ready(Some(Ok(Bytes::from(bytes))))
+                Poll::Ready(Some(Ok(self.take_complete_body())))
             }
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            Bytes::from(mem::take(self))
+        }
     }
 
     impl MessageBody for &'static str {
@@ -208,6 +343,14 @@ mod foreign_impls {
                 Poll::Ready(Some(Ok(bytes)))
             }
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            Bytes::from_static(mem::take(self).as_bytes())
+        }
     }
 
     impl MessageBody for String {
@@ -228,6 +371,14 @@ mod foreign_impls {
                 Poll::Ready(Some(Ok(Bytes::from(string))))
             }
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            Bytes::from(mem::take(self))
+        }
     }
 
     impl MessageBody for bytestring::ByteString {
@@ -244,6 +395,14 @@ mod foreign_impls {
             let string = mem::take(self.get_mut());
             Poll::Ready(Some(Ok(string.into_bytes())))
         }
+
+        fn is_complete_body(&self) -> bool {
+            true
+        }
+
+        fn take_complete_body(&mut self) -> Bytes {
+            mem::take(self).into_bytes()
+        }
     }
 }
 
@@ -406,6 +565,51 @@ mod tests {
         assert_poll_next!(pl, Bytes::from("test"));
     }
 
+    #[test]
+    fn take_string() {
+        let mut data = "test".repeat(2);
+        let data_bytes = Bytes::from(data.clone());
+        assert!(data.is_complete_body());
+        assert_eq!(data.take_complete_body(), data_bytes);
+
+        let mut big_data = "test".repeat(64 * 1024);
+        let data_bytes = Bytes::from(big_data.clone());
+        assert!(big_data.is_complete_body());
+        assert_eq!(big_data.take_complete_body(), data_bytes);
+    }
+
+    #[test]
+    fn take_boxed_equivalence() {
+        let mut data = Bytes::from_static(b"test");
+        assert!(data.is_complete_body());
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+
+        let mut data = Box::new(Bytes::from_static(b"test"));
+        assert!(data.is_complete_body());
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+
+        let mut data = Box::pin(Bytes::from_static(b"test"));
+        assert!(data.is_complete_body());
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+    }
+
+    #[test]
+    fn take_policy() {
+        let mut data = Bytes::from_static(b"test");
+        // first call returns chunk
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+        // second call returns empty
+        assert_eq!(data.take_complete_body(), b"".as_ref());
+
+        let waker = futures_util::task::noop_waker();
+        let mut cx = Context::from_waker(&waker);
+        let mut data = Bytes::from_static(b"test");
+        // take returns whole chunk
+        assert_eq!(data.take_complete_body(), b"test".as_ref());
+        // subsequent poll_next returns None
+        assert_eq!(Pin::new(&mut data).poll_next(&mut cx), Poll::Ready(None));
+    }
+
     // down-casting used to be done with a method on MessageBody trait
     // test is kept to demonstrate equivalence of Any trait
     #[actix_rt::test]

actix-http/src/body/none.rs

@@ -40,4 +40,14 @@ impl MessageBody for None {
     ) -> Poll<Option<Result<Bytes, Self::Error>>> {
         Poll::Ready(Option::None)
     }
+
+    #[inline]
+    fn is_complete_body(&self) -> bool {
+        true
+    }
+
+    #[inline]
+    fn take_complete_body(&mut self) -> Bytes {
+        Bytes::new()
+    }
 }