From acee48986699a6f0dd567b0f1176fcc32875ca45 Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Mon, 8 Apr 2024 00:27:54 +0300 Subject: refactor(tvix/nix-compat/wire): simplify write_bytes a bit 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 Autosubmit: flokli Tested-by: BuildkiteCI --- tvix/nix-compat/src/wire/bytes.rs | 19 ++++++++++++------- tvix/nix-compat/src/wire/bytes_writer.rs | 5 ++--- 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: &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], -- cgit 1.4.1