Repository: bearcove/loona (language: rust)

Issue:
WriteOwned and write
Hi. What's your take on whether or not writev should be interpreted as a write all?

I suspect most people would say it shouldn't be, as there is a count being returned.

So if an implementation wants to simulate a writev with a looping write, should it check the length of each write result and exit the loop early if the write wasn't complete?

And in the case of the owned API, should it have to rebuild the Vec<B> so all the entries are there for the caller to reuse, perhaps after modifying the vec slice and one of the vec entry slices?

Since the input vec argument ownership is passed in, it seems reasonable to replace the entries per write but I haven't tried to play with that yet to see how easy the ownership is kept straight, without resorting to unsafe.


Discussion / hints:
> What's your take on whether or not writev should be interpreted as a write all?

You're asking in the context of designing tokio-uring interfaces?

It seems the consensus (pre-uring) is to let writev do partial writes and "return" immediately, letting the caller "advance buffers" themselves. My first thought is that it'd be strange if writev did mutate the vec of buffers, since I don't think other writev implementations do that - giving ownership is "just" to make sure that the buffer doesn't vanish from under io_uring.

Perhaps that discussion is better suited to some issue on tokio-uring? You can tag me there if you want.
Sorry, I was too obtuse. I noticed your implementation of writev didn't allow for the kernel to do partial writes, each iteration was calling write, not write_all, and on error, it wasn't necessarily returning the full vec as it was building a new vec along the way while being successful.

I felt bad to bring it up because I know you haven't asked for help or bug reports and distractions like this are expensive and perhaps not worth it.
> I noticed your implementation of writev didn't allow for the kernel to do partial writes, each iteration was calling write, not write_all

Oh, nice catch!

> and on error, it wasn't necessarily returning the full vec as it was building a new vec along the way while being successful.

That one I don't see - in the `Err` arm, it does `out_list.extend(list)`, so it _should_ have all the input buffers?

Fix patch:
diff --git a/crates/hring-buffet/src/io.rs b/crates/hring-buffet/src/io.rs
index ed9556ad..7d98e7ff 100644
--- a/crates/hring-buffet/src/io.rs
+++ b/crates/hring-buffet/src/io.rs
@@ -13,18 +13,61 @@ pub trait ReadOwned {
 }
 
 pub trait WriteOwned {
+    /// Write a single buffer, taking ownership for the duration of the write.
+    /// Might perform a partial write, see [WriteOwned::write_all]
     async fn write<B: IoBuf>(&self, buf: B) -> BufResult<usize, B>;
 
+    /// Write a single buffer, re-trying the write if the kernel does a partial write.
+    async fn write_all<B: IoBuf>(&self, mut buf: B) -> std::io::Result<()> {
+        let mut written = 0;
+        let len = buf.bytes_init();
+        while written < len {
+            let (res, slice) = self.write(buf.slice(written..len)).await;
+            buf = slice.into_inner();
+            let n = res?;
+            if n == 0 {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::WriteZero,
+                    "write zero",
+                ));
+            }
+            written += n;
+        }
+        Ok(())
+    }
+
+    /// Write a list of buffers, taking ownership for the duration of the write.
+    /// Might perform a partial write, see [WriteOwned::writev_all]
     async fn writev<B: IoBuf>(&self, list: Vec<B>) -> BufResult<usize, Vec<B>> {
         let mut out_list = Vec::with_capacity(list.len());
         let mut list = list.into_iter();
         let mut total = 0;
 
         while let Some(buf) = list.next() {
+            let buf_len = buf.bytes_init();
             let (res, buf) = self.write(buf).await;
             out_list.push(buf);
+
             match res {
-                Ok(n) => total += n,
+                Ok(0) => {
+                    out_list.extend(list);
+                    return (
+                        Err(std::io::Error::new(
+                            std::io::ErrorKind::WriteZero,
+                            "write zero",
+                        )),
+                        out_list,
+                    );
+                }
+                Ok(n) => {
+                    total += n;
+                    if n < buf_len {
+                        // partial write, return the buffer list so the caller
+                        // might choose to try the write again
+                        out_list.extend(list);
+                        return (Ok(total), out_list);
+                    }
+                }
                 Err(e) => {
                     out_list.extend(list);
                     return (Err(e), out_list);
@@ -35,19 +78,107 @@ pub trait WriteOwned {
         (Ok(total), out_list)
     }
 
-    async fn write_all<B: IoBuf>(&self, mut buf: B) -> BufResult<(), B> {
-        let mut written = 0;
-        let len = buf.bytes_init();
-        while written < len {
-            let (res, slice) = self.write(buf.slice(written..len)).await;
-            buf = slice.into_inner();
-            let n = match res {
-                Ok(n) => n,
-                Err(e) => return (Err(e), buf),
-            };
-            written += n;
+    /// 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<()> {
+        // 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
+    }
+}
+
+/// Write a list of buffers, re-trying the write if the kernel does a partial write.
+async fn writev_all<B: IoBuf>(
+    this: &(impl WriteOwned + ?Sized),
+    list: Vec<B>,
+) -> std::io::Result<()> {
+    let mut list: Vec<_> = list.into_iter().map(BufOrSlice::Buf).collect();
+
+    while !list.is_empty() {
+        let res;
+        (res, list) = this.writev(list).await;
+        let n = res?;
+
+        if n == 0 {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::WriteZero,
+                "write zero",
+            ));
+        }
+
+        let mut n = n;
+        list = list
+            .into_iter()
+            .filter_map(|item| {
+                if n == 0 {
+                    Some(item)
+                } else {
+                    let item_len = item.len();
+
+                    if n >= item_len {
+                        n -= item_len;
+                        None
+                    } else {
+                        let item = item.consume(n);
+                        n = 0;
+                        Some(item)
+                    }
+                }
+            })
+            .collect();
+        assert_eq!(n, 0);
+    }
+
+    Ok(())
+}
+
+enum BufOrSlice<B: IoBuf> {
+    Buf(B),
+    Slice(tokio_uring::buf::Slice<B>),
+}
+
+unsafe impl<B: IoBuf> IoBuf for BufOrSlice<B> {
+    fn stable_ptr(&self) -> *const u8 {
+        match self {
+            BufOrSlice::Buf(b) => b.stable_ptr(),
+            BufOrSlice::Slice(s) => s.stable_ptr(),
+        }
+    }
+
+    fn bytes_init(&self) -> usize {
+        match self {
+            BufOrSlice::Buf(b) => b.bytes_init(),
+            BufOrSlice::Slice(s) => s.bytes_init(),
+        }
+    }
+
+    fn bytes_total(&self) -> usize {
+        match self {
+            BufOrSlice::Buf(b) => b.bytes_total(),
+            BufOrSlice::Slice(s) => s.bytes_total(),
+        }
+    }
+}
+
+impl<B: IoBuf> BufOrSlice<B> {
+    fn len(&self) -> usize {
+        match self {
+            BufOrSlice::Buf(b) => b.bytes_init(),
+            BufOrSlice::Slice(s) => s.len(),
+        }
+    }
+
+    /// Consume the first `n` bytes of the buffer (assuming they've been written).
+    /// This turns a `BufOrSlice::Buf` into a `BufOrSlice::Slice`
+    fn consume(self, n: usize) -> Self {
+        assert!(n <= self.len());
+
+        match self {
+            BufOrSlice::Buf(b) => BufOrSlice::Slice(b.slice(n..)),
+            BufOrSlice::Slice(s) => {
+                let n = s.begin() + n;
+                BufOrSlice::Slice(s.into_inner().slice(n..))
+            }
         }
-        (Ok(()), buf)
     }
 }
 
@@ -96,3 +227,81 @@ where
         self.1.write(buf).await
     }
 }
+
+#[cfg(all(test, not(feature = "miri")))]
+mod tests {
+    use std::{cell::RefCell, rc::Rc};
+
+    use crate::WriteOwned;
+
+    #[test]
+    fn test_write_all() {
+        enum Mode {
+            WriteZero,
+            WritePartial,
+        }
+
+        struct Writer {
+            mode: Mode,
+            bytes: Rc<RefCell<Vec<u8>>>,
+        }
+
+        impl WriteOwned for Writer {
+            async fn write<B: tokio_uring::buf::IoBuf>(
+                &self,
+                buf: B,
+            ) -> tokio_uring::BufResult<usize, B> {
+                assert!(buf.bytes_init() > 0, "zero-length writes are forbidden");
+
+                match self.mode {
+                    Mode::WriteZero => (Ok(0), buf),
+                    Mode::WritePartial => {
+                        let n = match buf.bytes_init() {
+                            1 => 1,
+                            _ => buf.bytes_init() / 2,
+                        };
+                        let slice = unsafe { std::slice::from_raw_parts(buf.stable_ptr(), n) };
+                        self.bytes.borrow_mut().extend_from_slice(slice);
+                        (Ok(n), buf)
+                    }
+                }
+            }
+        }
+
+        tokio_uring::start(async move {
+            let writer = Writer {
+                mode: Mode::WriteZero,
+                bytes: Default::default(),
+            };
+            let buf_a = vec![1, 2, 3, 4, 5];
+            let res = writer.write_all(buf_a).await;
+            assert!(res.is_err());
+
+            let writer = Writer {
+                mode: Mode::WriteZero,
+                bytes: Default::default(),
+            };
+            let buf_a = vec![1, 2, 3, 4, 5];
+            let buf_b = vec![6, 7, 8, 9, 10];
+            let res = writer.writev_all(vec![buf_a, buf_b]).await;
+            assert!(res.is_err());
+
+            let writer = Writer {
+                mode: Mode::WritePartial,
+                bytes: Default::default(),
+            };
+            let buf_a = vec![1, 2, 3, 4, 5];
+            writer.write_all(buf_a).await.unwrap();
+            assert_eq!(&writer.bytes.borrow()[..], &[1, 2, 3, 4, 5]);
+
+            let writer = Writer {
+                mode: Mode::WritePartial,
+                bytes: Default::default(),
+            };
+            let buf_a = vec![1, 2, 3, 4, 5];
+            let buf_b = vec![6, 7, 8, 9, 10];
+            writer.writev_all(vec![buf_a, buf_b]).await.unwrap();
+            assert_eq!(&writer.bytes.borrow()[..], &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
+        });
+    }
+}
diff --git a/crates/hring/src/h1/body.rs b/crates/hring/src/h1/body.rs
index a8261953..3d80cdbc 100644
--- a/crates/hring/src/h1/body.rs
+++ b/crates/hring/src/h1/body.rs
@@ -2,10 +2,7 @@ use std::{fmt, rc::Rc};
 
 use tracing::debug;
 
-use crate::{
-    util::{read_and_parse, write_all_list},
-    Body, BodyChunk, BodyErrorReason,
-};
+use crate::{util::read_and_parse, Body, BodyChunk, BodyErrorReason};
 use hring_buffet::{Piece, PieceList, ReadOwned, RollMut, WriteOwned};
 
 /// An HTTP/1.1 body, either chunked or content-length.
@@ -282,13 +279,10 @@ pub(crate) async fn write_h1_body_chunk(
             list.push(format!("{:x}\r\n", chunk.len()).into_bytes());
             list.push(chunk);
             list.push("\r\n");
-
-            let list = write_all_list(transport, list).await?;
-            drop(list);
+            transport.writev_all(list.into_vec()).await?;
         }
         BodyWriteMode::ContentLength => {
-            let (res, _) = transport.write_all(chunk).await;
-            res?;
+            transport.write_all(chunk).await?;
         }
         BodyWriteMode::Empty => {
             return Err(BodyErrorReason::CalledWriteBodyChunkWhenNoBodyWasExpected
@@ -306,9 +300,7 @@ pub(crate) async fn write_h1_body_end(
     debug!(?mode, "writing h1 body end");
     match mode {
         BodyWriteMode::Chunked => {
-            let mut list = PieceList::default();
-            list.push("0\r\n\r\n");
-            _ = write_all_list(transport, list).await?;
+            transport.write_all("0\r\n\r\n").await?;
         }
         BodyWriteMode::ContentLength => {
             // nothing to do
diff --git a/crates/hring/src/h1/client.rs b/crates/hring/src/h1/client.rs
index df820aa3..3bd82901 100644
--- a/crates/hring/src/h1/client.rs
+++ b/crates/hring/src/h1/client.rs
@@ -4,11 +4,7 @@ use eyre::Context;
 use http::header;
 use tracing::debug;
 
-use crate::{
-    types::Request,
-    util::{read_and_parse, write_all_list},
-    Body, HeadersExt, Response,
-};
+use crate::{types::Request, util::read_and_parse, Body, HeadersExt, Response};
 use hring_buffet::{PieceList, ReadWriteOwned, RollMut};
 
 use super::{
@@ -56,7 +52,8 @@ where
 
     let mut list = PieceList::default();
     encode_request(req, &mut list)?;
-    _ = write_all_list(transport.as_ref(), list)
+    transport
+        .writev_all(list.into_vec())
         .await
         .wrap_err("writing request headers")?;
 
diff --git a/crates/hring/src/h1/encode.rs b/crates/hring/src/h1/encode.rs
index 5db88644..181ea93c 100644
--- a/crates/hring/src/h1/encode.rs
+++ b/crates/hring/src/h1/encode.rs
@@ -5,7 +5,6 @@ use http::{StatusCode, Version};
 
 use crate::{
     types::{Headers, Request, Response},
-    util::write_all_list,
     Encoder,
 };
 use hring_buffet::{Piece, PieceList, WriteOwned};
@@ -152,13 +151,11 @@ where
         let mut list = PieceList::default();
         encode_response(res, &mut list)?;
 
-        let list = write_all_list(self.transport.as_ref(), list)
+        self.transport
+            .writev_all(list.into_vec())
             .await
             .wrap_err("writing response headers upstream")?;
 
-        // TODO: can we re-use that list? pool it?
-        drop(list);
-
         Ok(())
     }
 
@@ -178,13 +175,11 @@ where
         let mut list = PieceList::default();
         encode_headers(*trailers, &mut list)?;
 
-        let list = write_all_list(self.transport.as_ref(), list)
+        self.transport
+            .writev_all(list.into_vec())
             .await
             .wrap_err("writing response headers upstream")?;
 
-        // TODO: can we re-use that list? pool it?
-        drop(list);
-
         Ok(())
     }
 }
diff --git a/crates/hring/src/h1/server.rs b/crates/hring/src/h1/server.rs
index 97706912..b4173798 100644
--- a/crates/hring/src/h1/server.rs
+++ b/crates/hring/src/h1/server.rs
@@ -69,8 +69,10 @@ pub async fn serve(
             },
             Err(e) => {
                 if let Some(se) = e.downcast_ref::<SemanticError>() {
-                    let (res, _) = transport.write_all(se.as_http_response()).await;
-                    res.wrap_err("writing error response downstream")?;
+                    transport
+                        .write_all(se.as_http_response())
+                        .await
+                        .wrap_err("writing error response downstream")?;
                 }
 
                 debug!(?e, "error reading request header from downstream");
diff --git a/crates/hring/src/h2/parse.rs b/crates/hring/src/h2/parse.rs
index 55ed05e0..d9d77925 100644
--- a/crates/hring/src/h2/parse.rs
+++ b/crates/hring/src/h2/parse.rs
@@ -265,8 +265,7 @@ impl Frame {
             header.write_u32::<BigEndian>(self.stream_id.0)?;
         }
 
-        let (res, _) = w.write_all(header).await;
-        res?;
+        w.write_all(header).await?;
         Ok(())
     }
 }
diff --git a/crates/hring/src/h2/server.rs b/crates/hring/src/h2/server.rs
index ea3bb311..7cfa811a 100644
--- a/crates/hring/src/h2/server.rs
+++ b/crates/hring/src/h2/server.rs
@@ -545,8 +545,7 @@ async fn h2_write_loop(
                         frame.len = headers_encoded.len() as u32;
                         frame.write(transport.as_ref()).await?;
 
-                        let (res, _headers_encoded) = transport.write_all(headers_encoded).await;
-                        res?;
+                        transport.write_all(headers_encoded).await?;
                     }
                     H2EventPayload::BodyChunk(chunk) => {
                         let path = format!("/tmp/chunk-write-{index:06}.bin");
@@ -557,8 +556,7 @@ async fn h2_write_loop(
                         let frame = Frame::new(FrameType::Data(flags), ev.stream_id)
                             .with_len(chunk.len().try_into().unwrap());
                         frame.write(transport.as_ref()).await?;
-                        let (res, _) = transport.write_all(chunk).await;
-                        res?;
+                        transport.write_all(chunk).await?;
                     }
                     H2EventPayload::BodyEnd => {
                         let flags = DataFlags::EndStream;
@@ -573,8 +571,7 @@ async fn h2_write_loop(
                 let frame = Frame::new(FrameType::Ping(flags), StreamId::CONNECTION)
                     .with_len(payload.len() as u32);
                 frame.write(transport.as_ref()).await?;
-                let (res, _) = transport.write_all(payload).await;
-                res?;
+                transport.write_all(payload).await?;
             }
             H2ConnEvent::GoAway {
                 error_code,
@@ -597,11 +594,10 @@ async fn h2_write_loop(
                         .try_into()
                         .unwrap(),
                 );
+                // TODO: use `writev`
                 frame.write(transport.as_ref()).await?;
-                let (res, _) = transport.write_all(header).await;
-                res?;
-                let (res, _) = transport.write_all(additional_debug_data).await;
-                res?;
+                transport.write_all(header).await?;
+                transport.write_all(additional_debug_data).await?;
             }
         }
     }
diff --git a/crates/hring/src/util.rs b/crates/hring/src/util.rs
index 8025f85c..602c6030 100644
--- a/crates/hring/src/util.rs
+++ b/crates/hring/src/util.rs
@@ -3,7 +3,7 @@ use nom::IResult;
 use pretty_hex::PrettyHex;
 use tracing::{debug, trace};
 
-use hring_buffet::{PieceList, ReadOwned, Roll, RollMut, WriteOwned};
+use hring_buffet::{ReadOwned, Roll, RollMut};
 
 /// Returns `None` on EOF, error if partially parsed message.
 pub(crate) async fn read_and_parse<Parser, Output>(
@@ -76,28 +76,6 @@ where
     }
 }
 
-/// Write the filled part of a buffer to the given [TcpStream], returning a
-/// buffer re-using the remaining space.
-pub(crate) async fn write_all_list(
-    stream: &impl WriteOwned,
-    list: PieceList,
-) -> eyre::Result<PieceList> {
-    let len = list.len();
-    let num_chunks = list.num_pieces();
-    let list = list.into_vec();
-    debug!("writing {len} bytes in {num_chunks} chunks");
-
-    let (res, mut list) = stream.writev(list).await;
-    let n = res?;
-    debug!("wrote {n}/{len}");
-    if n < len {
-        unimplemented!();
-    }
-
-    list.clear();
-    Ok(list.into())
-}
-
 #[derive(thiserror::Error, Debug)]
 pub(crate) enum SemanticError {
     #[error("buffering limit reached while parsing")]


Test patch: