refactor(tvix/nix-compat/wire): simplify write_bytes a bit r/7872

Use the same EMPTY_BYTES trick from BytesWriter to write out the
padding, rather than allocating a Vec.

Change-Id: Ifb4ba1b45b7388adbc135fc8e46fd3d3cedd30aa
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11375
Reviewed-by: picnoir picnoir <picnoir@alternativebit.fr>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

2 files changed, 14 insertions, 10 deletions

diff --git a/tvix/nix-compat/src/wire/bytes.rs b/tvix/nix-compat/src/wire/bytes.rs
index e2f34ead6b..0a09305154 100644
--- a/tvix/nix-compat/src/wire/bytes.rs
+++ b/tvix/nix-compat/src/wire/bytes.rs
@@ -7,6 +7,9 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 use super::primitive;
 
+/// 8 null bytes, used to write out padding.
+pub(crate) const EMPTY_BYTES: &[u8; 8] = &[0u8; 8];
+
 #[allow(dead_code)]
 /// Read a "bytes wire packet" from the AsyncRead.
 /// Rejects reading more than `allowed_size` bytes of payload.
@@ -92,14 +95,16 @@ where
 /// [AsyncWriteExt] handle is buffered. This function is quite
 /// write-intesive.
 pub async fn write_bytes<W: AsyncWriteExt + Unpin>(w: &mut W, b: &[u8]) -> std::io::Result<()> {
-    // We're assuming the handle is buffered: we can afford not
-    // writing all the bytes in one go.
-    let len = b.len();
-    primitive::write_u64(w, len as u64).await?;
+    // write the size packet.
+    primitive::write_u64(w, b.len() as u64).await?;
+
+    // write the payload
     w.write_all(b).await?;
-    let padding = padding_len(len as u64);
-    if padding != 0 {
-        w.write_all(&vec![0; padding as usize]).await?;
+
+    // write padding if needed
+    let padding_len = padding_len(b.len() as u64) as usize;
+    if padding_len != 0 {
+        w.write_all(&EMPTY_BYTES[..padding_len]).await?;
     }
     Ok(())
 }
diff --git a/tvix/nix-compat/src/wire/bytes_writer.rs b/tvix/nix-compat/src/wire/bytes_writer.rs
index cf0e227e18..15cac91401 100644
--- a/tvix/nix-compat/src/wire/bytes_writer.rs
+++ b/tvix/nix-compat/src/wire/bytes_writer.rs
@@ -3,12 +3,11 @@ use std::task::{ready, Poll};
 
 use tokio::io::AsyncWrite;
 
+use super::bytes::EMPTY_BYTES;
+
 /// The length of the size field, in bytes is always 8.
 const LEN_SIZE: usize = 8;
 
-/// 8 null bytes, used to write out padding.
-const EMPTY_BYTES: &[u8; 8] = &[0u8; 8];
-
 pin_project! {
     /// Writes a "bytes wire packet" to the underlying writer.
     /// The format is the same as in [crate::wire::bytes::write_bytes],