From ab2b9e7644be39be675f2c79f898016fe13cc4bd Mon Sep 17 00:00:00 2001
From: Jonas Platte <jplatte+git@posteo.de>
Date: Fri, 9 Feb 2024 22:45:41 +0100
Subject: [PATCH] Optimize FromRequest for Bytes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

… and introduce FromRequest for BytesMut.

The more complex implementation is _currently_ faster, but will likely
be simplified to use http-body-util helpers again once the performance
of those catches up.

Co-authored-by: Yann Simon <yann.simon@commercetools.com>
---
 axum-core/src/extract/request_parts.rs | 79 +++++++++++++++++++++++---
 1 file changed, 72 insertions(+), 7 deletions(-)

diff --git a/axum-core/src/extract/request_parts.rs b/axum-core/src/extract/request_parts.rs
index 73f54db793..1ffb0ceac4 100644
--- a/axum-core/src/extract/request_parts.rs
+++ b/axum-core/src/extract/request_parts.rs
@@ -1,8 +1,9 @@
 use super::{rejection::*, FromRequest, FromRequestParts, Request};
 use crate::{body::Body, RequestExt};
 use async_trait::async_trait;
-use bytes::Bytes;
+use bytes::{Buf as _, BufMut, Bytes, BytesMut};
 use http::{request::Parts, Extensions, HeaderMap, Method, Uri, Version};
+use http_body::Body as _;
 use http_body_util::BodyExt;
 use std::convert::Infallible;
 
@@ -71,6 +72,37 @@ where
     }
 }
 
+#[async_trait]
+impl<S> FromRequest<S> for BytesMut
+where
+    S: Send + Sync,
+{
+    type Rejection = BytesRejection;
+
+    async fn from_request(req: Request, _: &S) -> Result<Self, Self::Rejection> {
+        let mut body = req.into_limited_body();
+        let mut bytes = BytesMut::new();
+        body_to_bytes_mut(&mut body, &mut bytes).await?;
+        Ok(bytes)
+    }
+}
+
+async fn body_to_bytes_mut(body: &mut Body, bytes: &mut BytesMut) -> Result<(), BytesRejection> {
+    while let Some(frame) = body
+        .frame()
+        .await
+        .transpose()
+        .map_err(FailedToBufferBody::from_err)?
+    {
+        let Ok(data) = frame.into_data() else {
+            return Ok(());
+        };
+        bytes.put(data);
+    }
+
+    Ok(())
+}
+
 #[async_trait]
 impl<S> FromRequest<S> for Bytes
 where
@@ -79,14 +111,47 @@ where
     type Rejection = BytesRejection;
 
     async fn from_request(req: Request, _: &S) -> Result<Self, Self::Rejection> {
-        let bytes = req
-            .into_limited_body()
-            .collect()
+        let mut body = req.into_limited_body();
+
+        // If there's only 1 chunk, we can just return Buf::to_bytes()
+        let first_chunk = if let Some(frame) = body
+            .frame()
             .await
+            .transpose()
             .map_err(FailedToBufferBody::from_err)?
-            .to_bytes();
-
-        Ok(bytes)
+        {
+            let Ok(first_chunk) = frame.into_data() else {
+                return Ok(Bytes::new());
+            };
+            first_chunk
+        } else {
+            return Ok(Bytes::new());
+        };
+
+        let mut bytes = if let Some(frame) = body
+            .frame()
+            .await
+            .transpose()
+            .map_err(FailedToBufferBody::from_err)?
+        {
+            let Ok(second_chunk) = frame.into_data() else {
+                return Ok(first_chunk);
+            };
+
+            let cap = first_chunk.remaining()
+                + second_chunk.remaining()
+                + body.size_hint().lower() as usize;
+
+            let mut bytes = BytesMut::with_capacity(cap);
+            bytes.put(first_chunk);
+            bytes.put(second_chunk);
+            bytes
+        } else {
+            return Ok(first_chunk);
+        };
+
+        body_to_bytes_mut(&mut body, &mut bytes).await?;
+        Ok(bytes.freeze())
     }
 }