Repository: bearcove/loona (language: rust)

Issue:
In read_and_parse, use `RollMut::reserve` rather than `RollMut::grow`
We don't necessarily need to double the capacity of the buffer, maybe we just got the end of a 4K buffer from previous parses, no need to go for a heap allocation there.


Discussion / hints:

Fix patch:
diff --git a/Justfile b/Justfile
index 09aa80a9..cfcfa069 100644
--- a/Justfile
+++ b/Justfile
@@ -33,7 +33,7 @@ bench *args:
 
 h2spec *args:
 	#!/bin/bash -eux
-	export RUST_LOG="${RUST_LOG:-debug,hpack=info}"
+	export RUST_LOG="${RUST_LOG:-hring=debug,hring_hpack=info}"
 	export RUST_BACKTRACE=1
 	cargo run --manifest-path test-crates/hring-h2spec/Cargo.toml -- {{args}}
 
diff --git a/crates/hring-buffet/src/roll.rs b/crates/hring-buffet/src/roll.rs
index 25ca11ff..46b3d919 100644
--- a/crates/hring-buffet/src/roll.rs
+++ b/crates/hring-buffet/src/roll.rs
@@ -402,6 +402,28 @@ impl RollMut {
     }
 }
 
+impl std::io::Write for RollMut {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        let n = buf.len();
+        if self.cap() < n {
+            // TODO: this is wrong, `reserve` might not reserve _enough data_.
+            // we know how much data we need here, so we should reserve exactly
+            // that amount (rounded up so it aligns nicely with the default
+            // buffer size)
+            self.reserve()
+                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+        }
+        self.put(buf)
+            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
+        Ok(n)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        // no need to flush
+        Ok(())
+    }
+}
+
 impl Deref for RollMut {
     type Target = [u8];
 
@@ -1327,4 +1349,13 @@ mod tests {
             buf.keep(rest);
         }
     }
+
+    #[test]
+    fn test_roll_io_write() {
+        let mut rm = RollMut::alloc().unwrap();
+        std::io::Write::write_all(&mut rm, b"hello").unwrap();
+
+        let roll = rm.take_all();
+        assert_eq!(std::str::from_utf8(&roll).unwrap(), "hello");
+    }
 }
diff --git a/crates/hring/Cargo.toml b/crates/hring/Cargo.toml
index eaf452cd..e53885e9 100644
--- a/crates/hring/Cargo.toml
+++ b/crates/hring/Cargo.toml
@@ -15,6 +15,7 @@ byteorder = "1.4.3"
 enum-repr = "0.2.6"
 enumflags2 = "0.7.5"
 eyre = { version = "0.6.8", default-features = false }
+futures-util = "0.3.25"
 hring-buffet = { version = "0.1.0", path = "../hring-buffet" }
 hring-hpack = { path = "../hring-hpack" }
 http = "0.2.8"
diff --git a/crates/hring/src/h1/body.rs b/crates/hring/src/h1/body.rs
index 3d80cdbc..babc74d4 100644
--- a/crates/hring/src/h1/body.rs
+++ b/crates/hring/src/h1/body.rs
@@ -275,11 +275,14 @@ pub(crate) async fn write_h1_body_chunk(
 ) -> eyre::Result<()> {
     match mode {
         BodyWriteMode::Chunked => {
-            let mut list = PieceList::default();
-            list.push(format!("{:x}\r\n", chunk.len()).into_bytes());
-            list.push(chunk);
-            list.push("\r\n");
-            transport.writev_all(list.into_vec()).await?;
+            transport
+                .writev_all(
+                    PieceList::default()
+                        .with(format!("{:x}\r\n", chunk.len()).into_bytes())
+                        .with(chunk)
+                        .with("\r\n"),
+                )
+                .await?;
         }
         BodyWriteMode::ContentLength => {
             transport.write_all(chunk).await?;
diff --git a/crates/hring/src/h1/client.rs b/crates/hring/src/h1/client.rs
index 3bd82901..10ae0fc1 100644
--- a/crates/hring/src/h1/client.rs
+++ b/crates/hring/src/h1/client.rs
@@ -50,10 +50,12 @@ where
         None => BodyWriteMode::Chunked,
     };
 
+    let mut buf = RollMut::alloc()?;
+
     let mut list = PieceList::default();
-    encode_request(req, &mut list)?;
+    encode_request(req, &mut list, &mut buf)?;
     transport
-        .writev_all(list.into_vec())
+        .writev_all(list)
         .await
         .wrap_err("writing request headers")?;
 
@@ -79,7 +81,6 @@ where
     let recv_res_fut = {
         let transport = transport.clone();
         async move {
-            let buf = RollMut::alloc()?;
             let (buf, res) = read_and_parse(
                 super::parse::response,
                 transport.as_ref(),
diff --git a/crates/hring/src/h1/encode.rs b/crates/hring/src/h1/encode.rs
index 181ea93c..b308062b 100644
--- a/crates/hring/src/h1/encode.rs
+++ b/crates/hring/src/h1/encode.rs
@@ -1,4 +1,4 @@
-use std::rc::Rc;
+use std::{io::Write, rc::Rc};
 
 use eyre::Context;
 use http::{StatusCode, Version};
@@ -7,15 +7,22 @@ use crate::{
     types::{Headers, Request, Response},
     Encoder,
 };
-use hring_buffet::{Piece, PieceList, WriteOwned};
+use hring_buffet::{Piece, PieceList, RollMut, WriteOwned};
 
 use super::body::{write_h1_body_chunk, write_h1_body_end, BodyWriteMode};
 
-pub(crate) fn encode_request(req: Request, list: &mut PieceList) -> eyre::Result<()> {
+pub(crate) fn encode_request(
+    req: Request,
+    list: &mut PieceList,
+    out_scratch: &mut RollMut,
+) -> eyre::Result<()> {
     list.push(req.method.into_chunk());
     list.push(" ");
-    // TODO: get rid of allocation here, or at least go through `RollMut`
-    list.push(req.uri.path().to_string().into_bytes());
+
+    assert_eq!(out_scratch.len(), 0);
+    out_scratch.write_all(req.uri.path().as_bytes())?;
+    list.push(out_scratch.take_all());
+
     match req.version {
         Version::HTTP_10 => list.push(" HTTP/1.0\r\n"),
         Version::HTTP_11 => list.push(" HTTP/1.1\r\n"),
@@ -152,7 +159,7 @@ where
         encode_response(res, &mut list)?;
 
         self.transport
-            .writev_all(list.into_vec())
+            .writev_all(list)
             .await
             .wrap_err("writing response headers upstream")?;
 
@@ -176,7 +183,7 @@ where
         encode_headers(*trailers, &mut list)?;
 
         self.transport
-            .writev_all(list.into_vec())
+            .writev_all(list)
             .await
             .wrap_err("writing response headers upstream")?;
 
diff --git a/crates/hring/src/h2/encode.rs b/crates/hring/src/h2/encode.rs
index c2566733..d7ba05f3 100644
--- a/crates/hring/src/h2/encode.rs
+++ b/crates/hring/src/h2/encode.rs
@@ -2,7 +2,7 @@ use std::fmt;
 
 use http::{StatusCode, Version};
 use tokio::sync::mpsc;
-use tracing::{debug, warn};
+use tracing::warn;
 
 use crate::{h1::body::BodyWriteMode, Encoder, Response};
 use hring_buffet::{Piece, Roll};
@@ -73,8 +73,6 @@ impl H2Encoder {
 
 impl Encoder for H2Encoder {
     async fn write_response(&mut self, res: Response) -> eyre::Result<()> {
-        debug!("H2Encoder::write_response");
-
         // TODO: don't panic here
         assert!(matches!(self.state, EncoderState::ExpectResponseHeaders));
 
@@ -90,8 +88,6 @@ impl Encoder for H2Encoder {
         chunk: hring_buffet::Piece,
         _mode: BodyWriteMode,
     ) -> eyre::Result<()> {
-        debug!("H2Encoder::write_body_chunk");
-
         assert!(matches!(self.state, EncoderState::ExpectResponseBody));
 
         self.send(H2EventPayload::BodyChunk(chunk)).await?;
@@ -100,8 +96,6 @@ impl Encoder for H2Encoder {
 
     // TODO: BodyWriteMode is not relevant for h2
     async fn write_body_end(&mut self, _mode: BodyWriteMode) -> eyre::Result<()> {
-        debug!("H2Encoder::write_body_end");
-
         assert!(matches!(self.state, EncoderState::ExpectResponseBody));
 
         self.send(H2EventPayload::BodyEnd).await?;
diff --git a/crates/hring/src/h2/parse.rs b/crates/hring/src/h2/parse.rs
index c55129cc..bb311e31 100644
--- a/crates/hring/src/h2/parse.rs
+++ b/crates/hring/src/h2/parse.rs
@@ -249,29 +249,22 @@ impl Frame {
         Ok((i, frame))
     }
 
-    pub const HEADER_LEN: usize = 9;
-
-    /// Serialize a frame header to the given buffer
-    pub fn serialize(&self, header: &mut [u8; Self::HEADER_LEN]) -> eyre::Result<()> {
-        let mut header = &mut header[..];
-
+    pub fn write_into(self, mut w: impl std::io::Write) -> eyre::Result<()> {
         use byteorder::{BigEndian, WriteBytesExt};
-        header.write_u24::<BigEndian>(self.len as _)?;
+        w.write_u24::<BigEndian>(self.len as _)?;
         let ft = self.frame_type.encode();
-        header.write_u8(ft.ty)?;
-        header.write_u8(ft.flags)?;
+        w.write_u8(ft.ty)?;
+        w.write_u8(ft.flags)?;
         // TODO: do we ever need to write the reserved bit?
-        header.write_u32::<BigEndian>(self.stream_id.0)?;
+        w.write_u32::<BigEndian>(self.stream_id.0)?;
 
         Ok(())
     }
 
-    /// Serialize a frame header to a `RollMut` whose filled portion is
-    /// empty, and return a `Roll` from it
-    pub fn to_roll(&self, buf: &mut RollMut) -> eyre::Result<Roll> {
-        buf.put_to_roll(Self::HEADER_LEN, |buf| {
-            self.serialize(buf.try_into().unwrap())
-        })
+    pub fn into_roll(self, mut scratch: &mut RollMut) -> eyre::Result<Roll> {
+        debug_assert_eq!(scratch.len(), 0);
+        self.write_into(&mut scratch)?;
+        Ok(scratch.take_all())
     }
 }
 
diff --git a/crates/hring/src/h2/server.rs b/crates/hring/src/h2/server.rs
index 59245adc..c087e090 100644
--- a/crates/hring/src/h2/server.rs
+++ b/crates/hring/src/h2/server.rs
@@ -1,6 +1,8 @@
 use std::{borrow::Cow, cell::RefCell, collections::HashMap, rc::Rc};
 
 use enumflags2::BitFlags;
+use eyre::Context;
+use futures_util::TryFutureExt;
 use http::{
     header::{self, HeaderName},
     uri::{Authority, PathAndQuery, Scheme},
@@ -106,7 +108,7 @@ pub async fn serve(
             StreamId::CONNECTION,
         );
         transport
-            .write_all(frame.to_roll(&mut out_scratch)?)
+            .write_all(frame.into_roll(&mut out_scratch)?)
             .await?;
         debug!("sent settings frame");
     }
@@ -120,13 +122,39 @@ pub async fn serve(
         client_buf,
         state.clone(),
     );
+
     let write_task = h2_write_loop(ev_rx, transport, out_scratch);
-    tokio::try_join!(read_task, write_task)?;
-    debug!("joined read_task / write_task");
+
+    let res = tokio::try_join!(
+        read_task.map_err(LoopError::Read),
+        write_task.map_err(LoopError::Write),
+    );
+    if let Err(e) = &res {
+        if let LoopError::Read(r) = e {
+            if r.downcast_ref::<ConnectionClosed>().is_some() {
+                return Ok(());
+            }
+        }
+        debug!("caught error from one of the tasks: {e} / {e:#?}");
+    }
+    res?;
 
     Ok(())
 }
 
+#[derive(thiserror::Error, Debug)]
+enum LoopError {
+    #[error("read error: {0}")]
+    Read(eyre::Report),
+
+    #[error("write error: {0}")]
+    Write(eyre::Report),
+}
+
+#[derive(thiserror::Error, Debug)]
+#[error("connection closed")]
+struct ConnectionClosed;
+
 async fn h2_read_loop(
     driver: Rc<impl ServerDriver + 'static>,
     ev_tx: mpsc::Sender<H2ConnEvent>,
@@ -144,11 +172,11 @@ async fn h2_read_loop(
                 Some((client_buf, frame)) => (client_buf, frame),
                 None => {
                     debug!("h2 client closed connection");
-                    return Ok(());
+                    return Err(eyre::Report::from(ConnectionClosed));
                 }
             };
 
-        debug!(?frame, "received h2 frame");
+        debug!(?frame, "Received");
 
         // TODO: there might be optimizations to be done for `Data` frames later
         // on, but for now, let's unconditionally read the payload (if it's not
@@ -235,7 +263,7 @@ async fn h2_read_loop(
                             let state = state.borrow();
                             if state.streams.contains_key(&frame.stream_id) {
                                 todo!(
-                                    "handle connection error: received headers for existing stream"
+                                    "handle connection error: received headers for existing stream {}", frame.stream_id
                                 );
                             }
                         }
@@ -271,7 +299,7 @@ async fn h2_read_loop(
                             (payload, _) = payload.split_at(at);
                         }
 
-                        debug!("receiving initial headers for stream {}", frame.stream_id);
+                        debug!("Receiving initial headers for stream {}", frame.stream_id);
                         let headers_data = HeadersData {
                             end_stream: flags.contains(HeadersFlags::EndStream),
                             fragments: smallvec![payload],
@@ -525,20 +553,23 @@ async fn h2_write_loop(
                     StreamId::CONNECTION,
                 );
                 transport
-                    .write_all(frame.to_roll(&mut out_scratch)?)
-                    .await?;
+                    .write_all(frame.into_roll(&mut out_scratch)?)
+                    .await
+                    .wrap_err("writing acknowledge settings")?;
             }
             H2ConnEvent::ServerEvent(ev) => {
-                debug!("Sending server event: {ev:?}");
+                debug!(?ev, "Writing");
 
                 match ev.payload {
                     H2EventPayload::Headers(res) => {
-                        debug!("Sending headers on stream {}", ev.stream_id);
                         let flags = HeadersFlags::EndHeaders;
                         let mut frame = Frame::new(FrameType::Headers(flags.into()), ev.stream_id);
 
-                        // TODO: don't allocate so much for headers
-                        // TODO: limt header size
+                        // TODO: don't allocate so much for headers. all `encode_into`
+                        // wants is an `IntoIter`, we can definitely have a custom iterator
+                        // that operates on all this instead of using a `Vec`.
+
+                        // TODO: limit header size
                         let mut headers: Vec<(&[u8], &[u8])> = vec![];
                         headers.push((b":status", res.status.as_str().as_bytes()));
                         for (name, value) in res.headers.iter() {
@@ -548,12 +579,18 @@ async fn h2_write_loop(
                             }
                             headers.push((name.as_str().as_bytes(), value));
                         }
-                        let headers_encoded = hpack_enc.encode(headers);
-                        frame.len = headers_encoded.len() as u32;
-                        let frame_roll = frame.to_roll(&mut out_scratch)?;
+
+                        assert_eq!(out_scratch.len(), 0);
+                        hpack_enc.encode_into(headers, &mut out_scratch)?;
+                        let fragment_block = out_scratch.take_all();
+
+                        frame.len = fragment_block.len() as u32;
+                        let frame_roll = frame.into_roll(&mut out_scratch)?;
+
                         transport
-                            .writev_all(PieceList::default().with(frame_roll).with(headers_encoded))
-                            .await?;
+                            .writev_all(PieceList::default().with(frame_roll).with(fragment_block))
+                            .await
+                            .wrap_err("writing headers")?;
                     }
                     H2EventPayload::BodyChunk(chunk) => {
                         let path = format!("/tmp/chunk-write-{index:06}.bin");
@@ -563,17 +600,19 @@ async fn h2_write_loop(
                         let flags = BitFlags::<DataFlags>::default();
                         let frame = Frame::new(FrameType::Data(flags), ev.stream_id)
                             .with_len(chunk.len().try_into().unwrap());
-                        let frame_roll = frame.to_roll(&mut out_scratch)?;
+                        let frame_roll = frame.into_roll(&mut out_scratch)?;
                         transport
                             .writev_all(PieceList::default().with(frame_roll).with(chunk))
-                            .await?;
+                            .await
+                            .wrap_err("writing bodychunk")?;
                     }
                     H2EventPayload::BodyEnd => {
                         let flags = DataFlags::EndStream;
                         let frame = Frame::new(FrameType::Data(flags.into()), ev.stream_id);
                         transport
-                            .write_all(frame.to_roll(&mut out_scratch)?)
-                            .await?;
+                            .write_all(frame.into_roll(&mut out_scratch)?)
+                            .await
+                            .wrap_err("writing BodyEnd")?;
                     }
                 }
             }
@@ -585,17 +624,18 @@ async fn h2_write_loop(
                 transport
                     .writev_all(
                         PieceList::default()
-                            .with(frame.to_roll(&mut out_scratch)?)
+                            .with(frame.into_roll(&mut out_scratch)?)
                             .with(payload),
                     )
-                    .await?;
+                    .await
+                    .wrap_err("writing pong")?;
             }
             H2ConnEvent::GoAway {
                 error_code,
                 last_stream_id,
                 additional_debug_data,
             } => {
-                debug!("must send goaway");
+                debug!(%last_stream_id, ?error_code, "Sending GoAway");
                 let header = out_scratch.put_to_roll(8, |mut slice| {
                     use byteorder::{BigEndian, WriteBytesExt};
                     // TODO: do we ever need to write the reserved bit?
@@ -605,7 +645,6 @@ async fn h2_write_loop(
                     Ok(())
                 })?;
 
-                debug!("sending goaway frame");
                 let frame = Frame::new(FrameType::GoAway, StreamId::CONNECTION).with_len(
                     (header.len() + additional_debug_data.len())
                         .try_into()
@@ -615,15 +654,15 @@ async fn h2_write_loop(
                 transport
                     .writev_all(
                         PieceList::default()
-                            .with(frame.to_roll(&mut out_scratch)?)
+                            .with(frame.into_roll(&mut out_scratch)?)
                             .with(header)
                             .with(additional_debug_data),
                     )
-                    .await?;
+                    .await
+                    .wrap_err("writing goaway")?;
             }
         }
     }
-    debug!("h2_write_loop is done");
 
     Ok(())
 }
@@ -776,7 +815,6 @@ fn end_headers(
         rx: piece_rx,
     };
 
-    debug!("Calling handler with the given body");
     tokio_uring::spawn({
         let driver = driver.clone();
         async move {
diff --git a/crates/hring/src/responder.rs b/crates/hring/src/responder.rs
index b62c0450..33c2eafc 100644
--- a/crates/hring/src/responder.rs
+++ b/crates/hring/src/responder.rs
@@ -1,5 +1,4 @@
 use http::header;
-use tracing::debug;
 
 use crate::{h1::body::BodyWriteMode, Body, BodyChunk, Headers, HeadersExt, Response};
 use hring_buffet::Piece;
@@ -86,8 +85,6 @@ where
         mut res: Response,
         body: &mut impl Body,
     ) -> eyre::Result<Responder<E, ResponseDone>> {
-        debug!("write final response with body");
-
         if let Some(clen) = body.content_len() {
             res.headers
                 .entry(header::CONTENT_LENGTH)
@@ -99,16 +96,13 @@ where
         }
 
         let mut this = self.write_final_response(res).await?;
-        debug!("wrote final response");
 
         loop {
             match body.next_chunk().await? {
                 BodyChunk::Chunk(chunk) => {
-                    debug!("proxying chunk of length {}", chunk.len());
                     this.write_chunk(chunk).await?;
                 }
                 BodyChunk::Done { trailers } => {
-                    debug!(mode = ?this.state.mode, "done proxying body");
                     // TODO: should we do something here in case of
                     // content-length mismatches?
                     return this.finish_body(trailers).await;
diff --git a/crates/hring/src/util.rs b/crates/hring/src/util.rs
index 602c6030..c2fb8877 100644
--- a/crates/hring/src/util.rs
+++ b/crates/hring/src/util.rs
@@ -44,8 +44,8 @@ where
                     }
 
                     if buf.cap() == 0 {
-                        trace!("buf had zero cap, growing");
-                        buf.grow()
+                        trace!("buf had zero cap, reserving");
+                        buf.reserve()?;
                     }
                     trace!(
                         "calling read_into, buf.cap={}, buf.len={} read_limit={read_limit}",


Test patch:
diff --git a/test-crates/hring-h2spec/Cargo.toml b/test-crates/hring-h2spec/Cargo.toml
index 08f9a57f..ee56bfb4 100644
--- a/test-crates/hring-h2spec/Cargo.toml
+++ b/test-crates/hring-h2spec/Cargo.toml
@@ -12,3 +12,6 @@ tokio = { version = "1.23.0", features = ["full"] }
 tracing = "0.1.37"
 tracing-subscriber = { version = "0.3.16", features = ["env-filter"] }
 which = "4.4.0"
+
+[profile.dev.package."*"]
+opt-level = 2
\ No newline at end of file
diff --git a/test-crates/hring-h2spec/src/main.rs b/test-crates/hring-h2spec/src/main.rs
index 5bf45b15..7a7e0405 100644
--- a/test-crates/hring-h2spec/src/main.rs
+++ b/test-crates/hring-h2spec/src/main.rs
@@ -125,6 +125,7 @@ async fn run_server(ln: TcpListener) -> color_eyre::Result<()> {
         let conf = Rc::new(hring::h2::ServerConf::default());
         let client_buf = RollMut::alloc()?;
         let driver = Rc::new(SDriver);
+
         tokio_uring::spawn(async move {
             if let Err(e) = hring::h2::serve(stream, conf, client_buf, driver).await {
                 tracing::error!("error serving client {}: {}", addr, e);