actix-http: play with allowing bytes passthrough in the h1 encoder

actix-http/examples/actix-web.rs

@@ -1,19 +1,34 @@
+use std::sync::OnceLock;
+
 use actix_http::HttpService;
 use actix_server::Server;
 use actix_service::map_config;
 use actix_web::{dev::AppConfig, get, App};
 
+static MEDIUM: OnceLock<String> = OnceLock::new();
+static LARGE: OnceLock<String> = OnceLock::new();
+
 #[get("/")]
 async fn index() -> &'static str {
     "Hello, world. From Actix Web!"
 }
 
+#[get("/large")]
+async fn large() -> &'static str {
+    LARGE.get_or_init(|| "123456890".repeat(1024 * 100))
+}
+
+#[get("/medium")]
+async fn medium() -> &'static str {
+    MEDIUM.get_or_init(|| "123456890".repeat(1024 * 5))
+}
+
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> std::io::Result<()> {
     Server::build()
         .bind("hello-world", "127.0.0.1:8080", || {
             // construct actix-web app
-            let app = App::new().service(index);
+            let app = App::new().service(index).service(large).service(medium);
 
             HttpService::build()
                 // pass the app to service builder

actix-http/src/h1/big_bytes.rs

@@ -0,0 +1,105 @@
+use std::collections::VecDeque;
+
+use bytes::{Buf, BufMut, Bytes, BytesMut};
+
+// 64KB max capacity (arbitrarily chosen)
+const MAX_CAPACITY: usize = 1024 * 64;
+
+pub(crate) struct BigBytes {
+    buffer: BytesMut,
+    frozen: VecDeque<Bytes>,
+    frozen_len: usize,
+}
+
+impl BigBytes {
+    pub(super) fn with_capacity(capacity: usize) -> Self {
+        Self {
+            buffer: BytesMut::with_capacity(capacity),
+            frozen: VecDeque::default(),
+            frozen_len: 0,
+        }
+    }
+
+    // Clear the internal queue and buffer, resetting length to zero
+    //
+    // if the internal buffer capacity exceeds 64KB or new_capacity, whichever is greater, it will
+    // be freed and a new buffer of capacity `new_capacity` will be allocated
+    pub(super) fn clear(&mut self, new_capacity: usize) {
+        std::mem::take(&mut self.frozen);
+        self.frozen_len = 0;
+        self.buffer.clear();
+
+        if self.buffer.capacity() > new_capacity.max(MAX_CAPACITY) {
+            self.buffer = BytesMut::with_capacity(new_capacity);
+        }
+    }
+
+    // Return a mutable reference to the underlying buffer. This should only be used when dealing
+    // with small allocations (e.g. writing headers)
+    pub(super) fn buffer_mut(&mut self) -> &mut BytesMut {
+        &mut self.buffer
+    }
+
+    pub(super) fn total_len(&mut self) -> usize {
+        self.frozen_len + self.buffer.len()
+    }
+
+    pub(super) fn is_empty(&self) -> bool {
+        self.frozen_len == 0 && self.buffer.is_empty()
+    }
+
+    // Add the `bytes` to the internal structure. If `bytes` exceeds 64KB, it is pushed into a
+    // queue, otherwise, it is added to a buffer.
+    pub(super) fn put_bytes(&mut self, bytes: Bytes) {
+        if !self.buffer.is_empty() {
+            let current = self.buffer.split().freeze();
+            self.frozen_len += current.len();
+            self.frozen.push_back(current);
+        }
+
+        if !bytes.is_empty() {
+            self.frozen_len += bytes.len();
+            self.frozen.push_back(bytes);
+        }
+    }
+
+    // Returns a slice of the frontmost buffer
+    pub(super) fn front_slice(&self) -> &[u8] {
+        if let Some(front) = self.frozen.front() {
+            front
+        } else {
+            &self.buffer
+        }
+    }
+
+    // Advances the first buffer by `count` bytes. If the first buffer is advanced to completion,
+    // it is popped from the queue
+    pub(super) fn advance(&mut self, count: usize) {
+        if let Some(front) = self.frozen.front_mut() {
+            front.advance(count);
+
+            if front.is_empty() {
+                self.frozen.pop_front();
+            }
+
+            self.frozen_len -= count;
+        } else {
+            self.buffer.advance(count);
+        }
+    }
+
+    // Drain the BibBytes, writing everything into the provided BytesMut
+    pub(super) fn write_to(&mut self, dst: &mut BytesMut) {
+        dst.reserve(self.total_len());
+
+        for buf in &self.frozen {
+            dst.put_slice(buf);
+        }
+
+        dst.put_slice(&self.buffer.split());
+
+        self.frozen_len = 0;
+
+        std::mem::take(&mut self.frozen);
+    }
+}

actix-http/src/h1/codec.rs

@@ -6,6 +6,7 @@ use http::{Method, Version};
 use tokio_util::codec::{Decoder, Encoder};
 
 use super::{
+    big_bytes::BigBytes,
     decoder::{self, PayloadDecoder, PayloadItem, PayloadType},
     encoder, Message, MessageType,
 };
@@ -146,14 +147,12 @@ impl Decoder for Codec {
     }
 }
 
-impl Encoder<Message<(Response<()>, BodySize)>> for Codec {
-    type Error = io::Error;
-
-    fn encode(
+impl Codec {
+    pub(super) fn encode_bigbytes(
         &mut self,
         item: Message<(Response<()>, BodySize)>,
-        dst: &mut BytesMut,
-    ) -> Result<(), Self::Error> {
+        dst: &mut BigBytes,
+    ) -> std::io::Result<()> {
         match item {
             Message::Item((mut res, length)) => {
                 // set response version
@@ -172,7 +171,7 @@ impl Encoder<Message<(Response<()>, BodySize)>> for Codec {
 
                 // encode message
                 self.encoder.encode(
-                    dst,
+                    dst.buffer_mut(),
                     &mut res,
                     self.flags.contains(Flags::HEAD),
                     self.flags.contains(Flags::STREAM),
@@ -184,18 +183,35 @@ impl Encoder<Message<(Response<()>, BodySize)>> for Codec {
             }
 
             Message::Chunk(Some(bytes)) => {
-                self.encoder.encode_chunk(bytes.as_ref(), dst)?;
+                self.encoder.encode_chunk_bigbytes(bytes, dst)?;
             }
 
             Message::Chunk(None) => {
-                self.encoder.encode_eof(dst)?;
+                self.encoder.encode_eof(dst.buffer_mut())?;
             }
         }
 
         Ok(())
     }
 }
 
+impl Encoder<Message<(Response<()>, BodySize)>> for Codec {
+    type Error = io::Error;
+
+    fn encode(
+        &mut self,
+        item: Message<(Response<()>, BodySize)>,
+        dst: &mut BytesMut,
+    ) -> Result<(), Self::Error> {
+        let mut bigbytes = BigBytes::with_capacity(1024 * 8);
+        self.encode_bigbytes(item, &mut bigbytes)?;
+
+        bigbytes.write_to(dst);
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

actix-http/src/h1/dispatcher.rs

@@ -12,14 +12,15 @@ use actix_codec::{Framed, FramedParts};
 use actix_rt::time::sleep_until;
 use actix_service::Service;
 use bitflags::bitflags;
-use bytes::{Buf, BytesMut};
+use bytes::BytesMut;
 use futures_core::ready;
 use pin_project_lite::pin_project;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tokio_util::codec::{Decoder as _, Encoder as _};
+use tokio_util::codec::Decoder as _;
 use tracing::{error, trace};
 
 use super::{
+    big_bytes::BigBytes,
     codec::Codec,
     decoder::MAX_BUFFER_SIZE,
     payload::{Payload, PayloadSender, PayloadStatus},
@@ -165,7 +166,7 @@ pin_project! {
 
         pub(super) io: Option<T>,
         read_buf: BytesMut,
-        write_buf: BytesMut,
+        write_buf: BigBytes,
         codec: Codec,
     }
 }
@@ -277,7 +278,7 @@ where
 
                     io: Some(io),
                     read_buf: BytesMut::with_capacity(HW_BUFFER_SIZE),
-                    write_buf: BytesMut::with_capacity(HW_BUFFER_SIZE),
+                    write_buf: BigBytes::with_capacity(HW_BUFFER_SIZE),
                     codec: Codec::new(config),
                 },
             },
@@ -329,27 +330,24 @@ where
         let InnerDispatcherProj { io, write_buf, .. } = self.project();
         let mut io = Pin::new(io.as_mut().unwrap());
 
-        let len = write_buf.len();
-        let mut written = 0;
-
-        while written < len {
-            match io.as_mut().poll_write(cx, &write_buf[written..])? {
+        while write_buf.total_len() > 0 {
+            match io.as_mut().poll_write(cx, write_buf.front_slice())? {
                 Poll::Ready(0) => {
+                    println!("WRITE ZERO");
                     error!("write zero; closing");
                     return Poll::Ready(Err(io::Error::new(io::ErrorKind::WriteZero, "")));
                 }
 
-                Poll::Ready(n) => written += n,
+                Poll::Ready(n) => write_buf.advance(n),
 
                 Poll::Pending => {
-                    write_buf.advance(written);
                     return Poll::Pending;
                 }
             }
         }
 
         // everything has written to I/O; clear buffer
-        write_buf.clear();
+        write_buf.clear(HW_BUFFER_SIZE);
 
         // flush the I/O and check if get blocked
         io.poll_flush(cx)
@@ -365,7 +363,7 @@ where
         let size = body.size();
 
         this.codec
-            .encode(Message::Item((res, size)), this.write_buf)
+            .encode_bigbytes(Message::Item((res, size)), this.write_buf)
             .map_err(|err| {
                 if let Some(mut payload) = this.payload.take() {
                     payload.set_error(PayloadError::Incomplete(None));
@@ -416,6 +414,7 @@ where
     fn send_continue(self: Pin<&mut Self>) {
         self.project()
             .write_buf
+            .buffer_mut()
             .extend_from_slice(b"HTTP/1.1 100 Continue\r\n\r\n");
     }
 
@@ -493,15 +492,16 @@ where
                 StateProj::SendPayload { mut body } => {
                     // keep populate writer buffer until buffer size limit hit,
                     // get blocked or finished.
-                    while this.write_buf.len() < super::payload::MAX_BUFFER_SIZE {
+                    while this.write_buf.total_len() < super::payload::MAX_BUFFER_SIZE {
                         match body.as_mut().poll_next(cx) {
                             Poll::Ready(Some(Ok(item))) => {
                                 this.codec
-                                    .encode(Message::Chunk(Some(item)), this.write_buf)?;
+                                    .encode_bigbytes(Message::Chunk(Some(item)), this.write_buf)?;
                             }
 
                             Poll::Ready(None) => {
-                                this.codec.encode(Message::Chunk(None), this.write_buf)?;
+                                this.codec
+                                    .encode_bigbytes(Message::Chunk(None), this.write_buf)?;
 
                                 // payload stream finished.
                                 // set state to None and handle next message
@@ -532,15 +532,16 @@ where
 
                     // keep populate writer buffer until buffer size limit hit,
                     // get blocked or finished.
-                    while this.write_buf.len() < super::payload::MAX_BUFFER_SIZE {
+                    while this.write_buf.total_len() < super::payload::MAX_BUFFER_SIZE {
                         match body.as_mut().poll_next(cx) {
                             Poll::Ready(Some(Ok(item))) => {
                                 this.codec
-                                    .encode(Message::Chunk(Some(item)), this.write_buf)?;
+                                    .encode_bigbytes(Message::Chunk(Some(item)), this.write_buf)?;
                             }
 
                             Poll::Ready(None) => {
-                                this.codec.encode(Message::Chunk(None), this.write_buf)?;
+                                this.codec
+                                    .encode_bigbytes(Message::Chunk(None), this.write_buf)?;
 
                                 // payload stream finished
                                 // set state to None and handle next message
@@ -575,6 +576,7 @@ where
                         // to service call.
                         Poll::Ready(Ok(req)) => {
                             this.write_buf
+                                .buffer_mut()
                                 .extend_from_slice(b"HTTP/1.1 100 Continue\r\n\r\n");
                             let fut = this.flow.service.call(req);
                             this.state.set(State::ServiceCall { fut });
@@ -1027,7 +1029,7 @@ where
             mem::take(this.codec),
             mem::take(this.read_buf),
         );
-        parts.write_buf = mem::take(this.write_buf);
+        this.write_buf.write_to(&mut parts.write_buf);
         let framed = Framed::from_parts(parts);
         this.flow.upgrade.as_ref().unwrap().call((req, framed))
     }