maurice/snix
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
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
This commit is contained in:
Florian Klink 2024-04-08 00:27:54 +03:00 • committed by clbot
parent eb910dfa3a
commit acee489866
2 changed files with 14 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

19
tvix/nix-compat/src/wire/bytes.rs
View file

@ -7,6 +7,9 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
use super::primitive; use super::primitive;
/// 8 null bytes, used to write out padding.
pub(crate) const EMPTY_BYTES: &[u8; 8] = &[0u8; 8];
#[allow(dead_code)] #[allow(dead_code)]
/// Read a "bytes wire packet" from the AsyncRead. /// Read a "bytes wire packet" from the AsyncRead.
/// Rejects reading more than `allowed_size` bytes of payload. /// Rejects reading more than `allowed_size` bytes of payload.
@ -92,14 +95,16 @@ where
/// [AsyncWriteExt] handle is buffered. This function is quite /// [AsyncWriteExt] handle is buffered. This function is quite
/// write-intesive. /// write-intesive.
pub async fn write_bytes<W: AsyncWriteExt + Unpin>(w: &mut W, b: &[u8]) -> std::io::Result<()> { 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 // write the size packet.
// writing all the bytes in one go. primitive::write_u64(w, b.len() as u64).await?;
let len = b.len();
primitive::write_u64(w, len as u64).await?; // write the payload
w.write_all(b).await?; w.write_all(b).await?;
let padding = padding_len(len as u64);
if padding != 0 { // write padding if needed
w.write_all(&vec![0; padding as usize]).await?; let padding_len = padding_len(b.len() as u64) as usize;
if padding_len != 0 {
w.write_all(&EMPTY_BYTES[..padding_len]).await?;
} }
Ok(()) Ok(())
} }

5
tvix/nix-compat/src/wire/bytes_writer.rs
View file

@ -3,12 +3,11 @@ use std::task::{ready, Poll};
use tokio::io::AsyncWrite; use tokio::io::AsyncWrite;
use super::bytes::EMPTY_BYTES;
/// The length of the size field, in bytes is always 8. /// The length of the size field, in bytes is always 8.
const LEN_SIZE: usize = 8; const LEN_SIZE: usize = 8;
/// 8 null bytes, used to write out padding.
const EMPTY_BYTES: &[u8; 8] = &[0u8; 8];
pin_project! { pin_project! {
/// Writes a "bytes wire packet" to the underlying writer. /// Writes a "bytes wire packet" to the underlying writer.
/// The format is the same as in [crate::wire::bytes::write_bytes], /// The format is the same as in [crate::wire::bytes::write_bytes],