From 6202c84408f8db554cf6d781d5ae21f5ce7ae6fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arpad=20M=C3=BCller?= <arpad@neon.tech>
Date: Tue, 2 Jul 2024 01:05:41 +0200
Subject: [PATCH] Actually return the same slice

---
 pageserver/src/virtual_file.rs | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/pageserver/src/virtual_file.rs b/pageserver/src/virtual_file.rs
index 869b7e6128..d671e70a9a 100644
--- a/pageserver/src/virtual_file.rs
+++ b/pageserver/src/virtual_file.rs
@@ -671,17 +671,27 @@ impl VirtualFile {
     }
 
     /// Writes `buf.slice(0..buf.bytes_init())`.
-    /// Returns the IoBuf that is underlying the BoundedBuf `buf`.
-    /// I.e., the returned value's `bytes_init()` method returns something different than the `bytes_init()` that was passed in.
-    /// It's quite brittle and easy to mis-use, so, we return the size in the Ok() variant.
+    ///
+    /// The returned `Slice<Buf>` is equivalent to the input `slice`, i.e., it's the same view into the same buffer.
+    /// We also return the amount of written data in the `Ok()` variant (although we could also have encountered
+    /// before an `Err()`).
     pub async fn write_all<Buf: IoBuf + Send>(
         &mut self,
         buf: Slice<Buf>,
         ctx: &RequestContext,
     ) -> (Slice<Buf>, Result<usize, Error>) {
+        let begin_end = (buf.begin(), buf.end());
+
+        macro_rules! return_ {
+            ($buf:expr, $val:expr) => {{
+                let buf = $buf.into_inner();
+                return (buf.slice(begin_end.0..begin_end.1), $val);
+            }};
+        }
+
         let nbytes = buf.bytes_init();
         if nbytes == 0 {
-            return (buf, Ok(0));
+            return_!(buf, Ok(0));
         }
         let mut buf = buf.slice(0..nbytes);
         while !buf.is_empty() {
@@ -689,22 +699,23 @@ impl VirtualFile {
             (buf, res) = self.write(buf, ctx).await;
             match res {
                 Ok(0) => {
-                    return (
+                    return_!(
                         buf,
                         Err(Error::new(
                             std::io::ErrorKind::WriteZero,
                             "failed to write whole buffer",
-                        )),
+                        ))
                     );
                 }
                 Ok(n) => {
                     buf = buf.slice(n..);
                 }
                 Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => {}
-                Err(e) => return (buf, Err(e)),
+                Err(e) => return_!(buf, Err(e)),
             }
         }
-        (buf, Ok(nbytes))
+
+        return_!(buf, Ok(nbytes));
     }
 
     async fn write<B: IoBuf + Send>(