Repository: bearcove/loona (language: rust)

Issue:
Use writev more heavily in h2 server, pool frame headers
`Frame::write` allocates a `Vec` of size 9 so we have a stable address for io_uring. We could allocate this once in the `h2_write_loop` and re-use it, since we only ever write one frame at a time.

Because `Frame::write` issues a `write_all` of its own, sometimes we call `write_all` as many as three times, when we could simply call `writev_all` once.


Discussion / hints:

Fix patch:
diff --git a/crates/hring-buffet/src/io.rs b/crates/hring-buffet/src/io.rs
index 7d98e7ff..50e79f11 100644
--- a/crates/hring-buffet/src/io.rs
+++ b/crates/hring-buffet/src/io.rs
@@ -79,10 +79,10 @@ pub trait WriteOwned {
     }
 
     /// Write a list of buffers, re-trying the write if the kernel does a partial write.
-    async fn writev_all<B: IoBuf>(&self, list: Vec<B>) -> std::io::Result<()> {
+    async fn writev_all<B: IoBuf>(&self, list: impl Into<Vec<B>>) -> std::io::Result<()> {
         // Workaround for https://github.com/rust-lang/rust/issues/107002,
         // remove after https://github.com/rust-lang/rust/pull/107013 is merged
-        writev_all(self, list).await
+        writev_all(self, list.into()).await
     }
 }
 
diff --git a/crates/hring-buffet/src/piece.rs b/crates/hring-buffet/src/piece.rs
index dfad7e19..68b999a7 100644
--- a/crates/hring-buffet/src/piece.rs
+++ b/crates/hring-buffet/src/piece.rs
@@ -147,6 +147,12 @@ impl PieceList {
         self.pieces.push(chunk.into());
     }
 
+    /// Add a single chunk to the list and return self
+    pub fn with(mut self, chunk: impl Into<Piece>) -> Self {
+        self.push(chunk);
+        self
+    }
+
     /// Returns total length
     pub fn len(&self) -> usize {
         self.pieces.iter().map(|c| c.len()).sum()
@@ -175,6 +181,12 @@ impl From<Vec<Piece>> for PieceList {
     }
 }
 
+impl From<PieceList> for Vec<Piece> {
+    fn from(list: PieceList) -> Self {
+        list.pieces
+    }
+}
+
 /// A piece of data with a stable address that's _also_
 /// valid utf-8.
 #[derive(Clone)]
diff --git a/crates/hring-buffet/src/roll.rs b/crates/hring-buffet/src/roll.rs
index c5f5b1ba..25ca11ff 100644
--- a/crates/hring-buffet/src/roll.rs
+++ b/crates/hring-buffet/src/roll.rs
@@ -274,6 +274,47 @@ impl RollMut {
         Ok(())
     }
 
+    /// Put data into this RollMut with a closure. Panics if `len > self.cap()`
+    pub fn put_with<T, E>(
+        &mut self,
+        len: usize,
+        f: impl FnOnce(&mut [u8]) -> Result<T, E>,
+    ) -> Result<T, E> {
+        assert!(len <= self.cap());
+
+        let u32_len: u32 = len.try_into().unwrap();
+        let slice = unsafe {
+            std::slice::from_raw_parts_mut(self.storage.as_mut_ptr().add(self.len as usize), len)
+        };
+        let res = f(slice);
+        if res.is_ok() {
+            self.len += u32_len;
+        }
+        res
+    }
+
+    /// Assert that this RollMut isn't filled at all, reserve enough size to put
+    /// `len` bytes in it, then fill it with the given closure. If the closure
+    /// returns `Err`, it's as if the put never happened.
+    pub fn put_to_roll(
+        &mut self,
+        len: usize,
+        f: impl FnOnce(&mut [u8]) -> eyre::Result<()>,
+    ) -> eyre::Result<Roll> {
+        // TODO: this whole dance is a bit silly: the idea is to have a
+        // `RollMut` around that we use whenever we need to serialize something.
+        // it's weird that we need to do all this for it to happen but ah well.
+
+        assert_eq!(self.len(), 0);
+        if self.cap() < len {
+            self.reserve()?;
+        }
+        self.put_with(len, f)?;
+        let roll = self.take_all();
+        debug_assert_eq!(roll.len(), len);
+        Ok(roll)
+    }
+
     /// Get a [Roll] corresponding to the filled portion of this buffer
     pub fn filled(&self) -> Roll {
         match &self.storage {
diff --git a/crates/hring/src/h2/parse.rs b/crates/hring/src/h2/parse.rs
index d9d77925..c55129cc 100644
--- a/crates/hring/src/h2/parse.rs
+++ b/crates/hring/src/h2/parse.rs
@@ -14,7 +14,7 @@ use nom::{
     IResult,
 };
 
-use hring_buffet::{Roll, WriteOwned};
+use hring_buffet::{Roll, RollMut};
 
 /// This is sent by h2 clients after negotiating over ALPN, or when doing h2c.
 pub const PREFACE: &[u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";
@@ -249,25 +249,30 @@ impl Frame {
         Ok((i, frame))
     }
 
-    /// Write a frame (without payload)
-    pub async fn write(&self, w: &impl WriteOwned) -> eyre::Result<()> {
-        // TODO: we need an owned type here, but do we need to re-allocate it
-        // here? can we have a pool of them?
-        let mut header = vec![0u8; 9];
-        {
-            use byteorder::{BigEndian, WriteBytesExt};
-            let mut header = &mut header[..];
-            header.write_u24::<BigEndian>(self.len as _)?;
-            let ft = self.frame_type.encode();
-            header.write_u8(ft.ty)?;
-            header.write_u8(ft.flags)?;
-            // TODO: do we ever need to write the reserved bit?
-            header.write_u32::<BigEndian>(self.stream_id.0)?;
-        }
+    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[..];
+
+        use byteorder::{BigEndian, WriteBytesExt};
+        header.write_u24::<BigEndian>(self.len as _)?;
+        let ft = self.frame_type.encode();
+        header.write_u8(ft.ty)?;
+        header.write_u8(ft.flags)?;
+        // TODO: do we ever need to write the reserved bit?
+        header.write_u32::<BigEndian>(self.stream_id.0)?;
 
-        w.write_all(header).await?;
         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())
+        })
+    }
 }
 
 /// See https://httpwg.org/specs/rfc9113.html#FrameHeader - the first bit
diff --git a/crates/hring/src/h2/server.rs b/crates/hring/src/h2/server.rs
index 7cfa811a..59245adc 100644
--- a/crates/hring/src/h2/server.rs
+++ b/crates/hring/src/h2/server.rs
@@ -23,7 +23,7 @@ use crate::{
     util::read_and_parse,
     ExpectResponseHeaders, Headers, Method, Request, Responder, ServerDriver,
 };
-use hring_buffet::{Piece, PieceStr, ReadWriteOwned, Roll, RollMut};
+use hring_buffet::{Piece, PieceList, PieceStr, ReadWriteOwned, Roll, RollMut};
 
 /// HTTP/2 server configuration
 pub struct ServerConf {
@@ -97,13 +97,17 @@ pub async fn serve(
     };
     debug!("read preface");
 
+    let mut out_scratch = RollMut::alloc()?;
+
     // we have to send a settings frame
     {
         let frame = Frame::new(
             FrameType::Settings(Default::default()),
             StreamId::CONNECTION,
         );
-        frame.write(transport.as_ref()).await?;
+        transport
+            .write_all(frame.to_roll(&mut out_scratch)?)
+            .await?;
         debug!("sent settings frame");
     }
 
@@ -116,7 +120,7 @@ pub async fn serve(
         client_buf,
         state.clone(),
     );
-    let write_task = h2_write_loop(ev_rx, transport);
+    let write_task = h2_write_loop(ev_rx, transport, out_scratch);
     tokio::try_join!(read_task, write_task)?;
     debug!("joined read_task / write_task");
 
@@ -505,6 +509,7 @@ async fn send_goaway(
 async fn h2_write_loop(
     mut ev_rx: mpsc::Receiver<H2ConnEvent>,
     transport: Rc<impl ReadWriteOwned>,
+    mut out_scratch: RollMut,
 ) -> eyre::Result<()> {
     let mut hpack_enc = hring_hpack::Encoder::new();
 
@@ -515,11 +520,13 @@ async fn h2_write_loop(
         match ev {
             H2ConnEvent::AcknowledgeSettings => {
                 debug!("acknowledging new settings");
-                let res_frame = Frame::new(
+                let frame = Frame::new(
                     FrameType::Settings(SettingsFlags::Ack.into()),
                     StreamId::CONNECTION,
                 );
-                res_frame.write(transport.as_ref()).await?;
+                transport
+                    .write_all(frame.to_roll(&mut out_scratch)?)
+                    .await?;
             }
             H2ConnEvent::ServerEvent(ev) => {
                 debug!("Sending server event: {ev:?}");
@@ -543,9 +550,10 @@ async fn h2_write_loop(
                         }
                         let headers_encoded = hpack_enc.encode(headers);
                         frame.len = headers_encoded.len() as u32;
-                        frame.write(transport.as_ref()).await?;
-
-                        transport.write_all(headers_encoded).await?;
+                        let frame_roll = frame.to_roll(&mut out_scratch)?;
+                        transport
+                            .writev_all(PieceList::default().with(frame_roll).with(headers_encoded))
+                            .await?;
                     }
                     H2EventPayload::BodyChunk(chunk) => {
                         let path = format!("/tmp/chunk-write-{index:06}.bin");
@@ -555,13 +563,17 @@ 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());
-                        frame.write(transport.as_ref()).await?;
-                        transport.write_all(chunk).await?;
+                        let frame_roll = frame.to_roll(&mut out_scratch)?;
+                        transport
+                            .writev_all(PieceList::default().with(frame_roll).with(chunk))
+                            .await?;
                     }
                     H2EventPayload::BodyEnd => {
                         let flags = DataFlags::EndStream;
                         let frame = Frame::new(FrameType::Data(flags.into()), ev.stream_id);
-                        frame.write(transport.as_ref()).await?;
+                        transport
+                            .write_all(frame.to_roll(&mut out_scratch)?)
+                            .await?;
                     }
                 }
             }
@@ -570,8 +582,13 @@ async fn h2_write_loop(
                 let flags = PingFlags::Ack.into();
                 let frame = Frame::new(FrameType::Ping(flags), StreamId::CONNECTION)
                     .with_len(payload.len() as u32);
-                frame.write(transport.as_ref()).await?;
-                transport.write_all(payload).await?;
+                transport
+                    .writev_all(
+                        PieceList::default()
+                            .with(frame.to_roll(&mut out_scratch)?)
+                            .with(payload),
+                    )
+                    .await?;
             }
             H2ConnEvent::GoAway {
                 error_code,
@@ -579,14 +596,14 @@ async fn h2_write_loop(
                 additional_debug_data,
             } => {
                 debug!("must send goaway");
-                let mut header = vec![0u8; 8];
-                {
+                let header = out_scratch.put_to_roll(8, |mut slice| {
                     use byteorder::{BigEndian, WriteBytesExt};
-                    let mut header = &mut header[..];
                     // TODO: do we ever need to write the reserved bit?
-                    header.write_u32::<BigEndian>(last_stream_id.0)?;
-                    header.write_u32::<BigEndian>(error_code.repr())?;
-                }
+                    slice.write_u32::<BigEndian>(last_stream_id.0)?;
+                    slice.write_u32::<BigEndian>(error_code.repr())?;
+
+                    Ok(())
+                })?;
 
                 debug!("sending goaway frame");
                 let frame = Frame::new(FrameType::GoAway, StreamId::CONNECTION).with_len(
@@ -594,10 +611,15 @@ async fn h2_write_loop(
                         .try_into()
                         .unwrap(),
                 );
-                // TODO: use `writev`
-                frame.write(transport.as_ref()).await?;
-                transport.write_all(header).await?;
-                transport.write_all(additional_debug_data).await?;
+
+                transport
+                    .writev_all(
+                        PieceList::default()
+                            .with(frame.to_roll(&mut out_scratch)?)
+                            .with(header)
+                            .with(additional_debug_data),
+                    )
+                    .await?;
             }
         }
     }


Test patch: