mdebray/tvl-depot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(tvix): Avoid buffering file into memory in builtins.hashFile

Browse source 
Right now `builtins.hashFile` always reads the entire file into memory
before hashing, which is not ideal for large files. This replaces
`read_to_string` with `open_file` which allows calculating the hash of
the file without buffering it entirely into memory. Other callers can
continue to buffer into memory if they choose, but they still use the
`open_file` VM request and then call `read_to_string` or `read_to_end`
on the `std::io::Reader`.

Fixes b/380

Change-Id: Ifa1c8324bcee8f751604b0b449feab875c632fda
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11236
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
This commit is contained in:
Connor Brewster 2024-03-22 18:52:21 -05:00
parent 17849c5c00
commit 63116d8c21
9 changed files with 80 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
tvix/eval/src/builtins/hash.rs
View file

@ -6,18 +6,22 @@ use sha2::{digest::Output, Digest, Sha256, Sha512};
use crate::ErrorKind;
fn hash<D: Digest>(b: &[u8]) -> Output<D> {
/// Reads through all data from the passed reader, and returns the resulting [Digest].
/// The exact hash function used is left generic over all [Digest].
fn hash<D: Digest + std::io::Write>(mut r: impl std::io::Read) -> Result<Output<D>, ErrorKind> {
let mut hasher = D::new();
hasher.update(b);
hasher.finalize()
std::io::copy(&mut r, &mut hasher)?;
Ok(hasher.finalize())
}
pub fn hash_nix_string(algo: impl AsRef<[u8]>, s: impl AsRef<[u8]>) -> Result<String, ErrorKind> {
/// For a given algo "string" and reader for data, calculate the digest
/// and return it as a hexlower encoded [String].
pub fn hash_nix_string(algo: impl AsRef<[u8]>, s: impl std::io::Read) -> Result<String, ErrorKind> {
match algo.as_ref() {
b"md5" => Ok(HEXLOWER.encode(hash::<Md5>(s.as_ref()).as_bstr())),
b"sha1" => Ok(HEXLOWER.encode(hash::<Sha1>(s.as_ref()).as_bstr())),
b"sha256" => Ok(HEXLOWER.encode(hash::<Sha256>(s.as_ref()).as_bstr())),
b"sha512" => Ok(HEXLOWER.encode(hash::<Sha512>(s.as_ref()).as_bstr())),
b"md5" => Ok(HEXLOWER.encode(hash::<Md5>(s)?.as_bstr())),
b"sha1" => Ok(HEXLOWER.encode(hash::<Sha1>(s)?.as_bstr())),
b"sha256" => Ok(HEXLOWER.encode(hash::<Sha256>(s)?.as_bstr())),
b"sha512" => Ok(HEXLOWER.encode(hash::<Sha512>(s)?.as_bstr())),
_ => Err(ErrorKind::UnknownHashType(
algo.as_ref().as_bstr().to_string(),
)),

15
tvix/eval/src/builtins/impure.rs
View file

@ -31,14 +31,13 @@ mod impure_builtins {
}
#[builtin("hashFile")]
#[allow(non_snake_case)]
async fn builtin_hashFile(co: GenCo, algo: Value, path: Value) -> Result<Value, ErrorKind> {
async fn builtin_hash_file(co: GenCo, algo: Value, path: Value) -> Result<Value, ErrorKind> {
let path = match coerce_value_to_path(&co, path).await? {
Err(cek) => return Ok(Value::from(cek)),
Ok(p) => p,
};
let s = generators::request_read_to_string(&co, path).await;
hash_nix_string(algo.to_str()?, s.to_str()?).map(Value::from)
let r = generators::request_open_file(&co, path).await;
Ok(hash_nix_string(algo.to_str()?, r).map(Value::from)?)
}
#[builtin("pathExists")]
@ -79,7 +78,13 @@ mod impure_builtins {
async fn builtin_read_file(co: GenCo, path: Value) -> Result<Value, ErrorKind> {
match coerce_value_to_path(&co, path).await? {
Err(cek) => Ok(Value::from(cek)),
Ok(path) => Ok(generators::request_read_to_string(&co, path).await),
Ok(path) => {
let mut buf = Vec::new();
generators::request_open_file(&co, path)
.await
.read_to_end(&mut buf)?;
Ok(Value::from(buf))
}
}
}
}

5
tvix/eval/src/builtins/mod.rs
View file

@ -773,9 +773,8 @@ mod pure_builtins {
}
#[builtin("hashString")]
#[allow(non_snake_case)]
async fn builtin_hashString(co: GenCo, algo: Value, s: Value) -> Result<Value, ErrorKind> {
hash_nix_string(algo.to_str()?, s.to_str()?).map(Value::from)
async fn builtin_hash_string(co: GenCo, algo: Value, s: Value) -> Result<Value, ErrorKind> {
hash_nix_string(algo.to_str()?, std::io::Cursor::new(s.to_str()?)).map(Value::from)
}
#[builtin("head")]

9
tvix/eval/src/compiler/import.rs
View file

@ -6,7 +6,6 @@
//! instance, or observers).
use super::GlobalsMap;
use bstr::ByteSlice;
use genawaiter::rc::Gen;
use std::rc::Weak;
@ -39,9 +38,11 @@ async fn import_impl(
return Ok(cached);
}
// TODO(tazjin): make this return a string directly instead
let contents: Value = generators::request_read_to_string(&co, path.clone()).await;
let contents = contents.to_str()?.to_str()?.to_owned();
let mut reader = generators::request_open_file(&co, path.clone()).await;
// We read to a String instead of a Vec<u8> because rnix only supports
// string source files.
let mut contents = String::new();
reader.read_to_string(&mut contents)?;
let parsed = rnix::ast::Root::parse(&contents);
let errors = parsed.errors();

16
tvix/eval/src/io.rs
View file

@ -16,6 +16,7 @@
//! how store paths are opened and so on.
use std::{
fs::File,
io,
path::{Path, PathBuf},
};
@ -48,13 +49,8 @@ pub trait EvalIO {
/// * `builtins.pathExists :: path -> bool`
fn path_exists(&self, path: &Path) -> io::Result<bool>;
/// Read the file at the specified path to a `Vec<u8>`.
///
/// This is used for the following language evaluation cases:
///
/// * `builtins.readFile :: path -> string`
/// * `builtins.import :: path -> any`
fn read_to_end(&self, path: &Path) -> io::Result<Vec<u8>>;
/// Open the file at the specified path to a `io::Read`.
fn open(&self, path: &Path) -> io::Result<Box<dyn io::Read>>;
/// Read the directory at the specified path and return the names
/// of its entries associated with their [`FileType`].
@ -99,8 +95,8 @@ impl EvalIO for StdIO {
path.try_exists()
}
fn read_to_end(&self, path: &Path) -> io::Result<Vec<u8>> {
std::fs::read(path)
fn open(&self, path: &Path) -> io::Result<Box<dyn io::Read>> {
Ok(Box::new(File::open(path)?))
}
fn read_dir(&self, path: &Path) -> io::Result<Vec<(bytes::Bytes, FileType)>> {
@ -145,7 +141,7 @@ impl EvalIO for DummyIO {
))
}
fn read_to_end(&self, _: &Path) -> io::Result<Vec<u8>> {
fn open(&self, _: &Path) -> io::Result<Box<dyn io::Read>> {
Err(io::Error::new(
io::ErrorKind::Unsupported,
"I/O methods are not implemented in DummyIO",

27
tvix/eval/src/vm/generators.rs
View file

@ -102,8 +102,8 @@ pub enum VMRequest {
/// Request that the VM imports the given path through its I/O interface.
PathImport(PathBuf),
/// Request that the VM reads the given path to a string.
ReadToString(PathBuf),
/// Request that the VM opens the specified file and provides a reader.
OpenFile(PathBuf),
/// Request that the VM checks whether the given path exists.
PathExists(PathBuf),
@ -170,8 +170,8 @@ impl Display for VMRequest {
write!(f, "import_cache_put({})", p.to_string_lossy())
}
VMRequest::PathImport(p) => write!(f, "path_import({})", p.to_string_lossy()),
VMRequest::ReadToString(p) => {
write!(f, "read_to_string({})", p.to_string_lossy())
VMRequest::OpenFile(p) => {
write!(f, "open_file({})", p.to_string_lossy())
}
VMRequest::PathExists(p) => write!(f, "path_exists({})", p.to_string_lossy()),
VMRequest::ReadDir(p) => write!(f, "read_dir({})", p.to_string_lossy()),
@ -199,6 +199,9 @@ pub enum VMResponse {
/// VM response with a span to use at the current point.
Span(LightSpan),
/// [std::io::Reader] produced by the VM in response to some IO operation.
Reader(Box<dyn std::io::Read>),
}
impl Display for VMResponse {
@ -209,6 +212,7 @@ impl Display for VMResponse {
VMResponse::Path(p) => write!(f, "path({})", p.to_string_lossy()),
VMResponse::Directory(d) => write!(f, "dir(len = {})", d.len()),
VMResponse::Span(_) => write!(f, "span"),
VMResponse::Reader(_) => write!(f, "reader"),
}
}
}
@ -425,18 +429,18 @@ where
message = VMResponse::Path(imported);
}
VMRequest::ReadToString(path) => {
let content = self
VMRequest::OpenFile(path) => {
let reader = self
.io_handle
.as_ref()
.read_to_end(&path)
.open(&path)
.map_err(|e| ErrorKind::IO {
path: Some(path),
error: e.into(),
})
.with_span(&span, self)?;
message = VMResponse::Value(content.into())
message = VMResponse::Reader(reader)
}
VMRequest::PathExists(path) => {
@ -730,9 +734,10 @@ pub(crate) async fn request_path_import(co: &GenCo, path: PathBuf) -> PathBuf {
}
}
pub(crate) async fn request_read_to_string(co: &GenCo, path: PathBuf) -> Value {
match co.yield_(VMRequest::ReadToString(path)).await {
VMResponse::Value(value) => value,
/// Request that the VM open a [std::io::Read] for the specified file.
pub async fn request_open_file(co: &GenCo, path: PathBuf) -> Box<dyn std::io::Read> {
match co.yield_(VMRequest::OpenFile(path)).await {
VMResponse::Reader(value) => value,
msg => panic!(
"Tvix bug: VM responded with incorrect generator message: {}",
msg

6
tvix/glue/src/builtins/import.rs
View file

@ -177,9 +177,9 @@ mod import_builtins {
})
.transpose()?;
// FUTUREWORK(performance): this reads the file instead of using a stat-like
// system call to the file, this degrades very badly on large files.
if !recursive_ingestion && state.read_to_end(path.as_ref()).is_err() {
// FUTUREWORK(performance): this opens the file instead of using a stat-like
// system call to the file.
if !recursive_ingestion && state.open(path.as_ref()).is_err() {
Err(ImportError::FlatImportOfNonFile(
path.to_string_lossy().to_string(),
))?;

11
tvix/glue/src/tvix_io.rs
View file

@ -8,7 +8,7 @@
//! otherwise fundamental features like nixpkgs bootstrapping and hash
//! calculation will not work.
use std::io;
use std::io::{self, Cursor};
use std::path::{Path, PathBuf};
use tvix_eval::{EvalIO, FileType};
@ -44,7 +44,7 @@ where
self.actual.as_ref().path_exists(path)
}
fn read_to_end(&self, path: &Path) -> io::Result<Vec<u8>> {
fn open(&self, path: &Path) -> io::Result<Box<dyn io::Read>> {
// Bundled version of corepkgs/fetchurl.nix. The counterpart
// of this happens in [crate::configure_nix_path], where the `nix_path`
// of the evaluation has `nix=/__corepkgs__` added to it.
@ -52,13 +52,12 @@ where
// This workaround is similar to what cppnix does for passing
// the path through.
//
// TODO: this comparison is bad and allocates, we should use
// the sane path library.
// TODO: this comparison is bad we should use the sane path library.
if path.starts_with("/__corepkgs__/fetchurl.nix") {
return Ok(include_bytes!("fetchurl.nix").to_vec());
return Ok(Box::new(Cursor::new(include_bytes!("fetchurl.nix"))));
}
self.actual.as_ref().read_to_end(path)
self.actual.as_ref().open(path)
}
fn read_dir(&self, path: &Path) -> io::Result<Vec<(bytes::Bytes, FileType)>> {

23
tvix/glue/src/tvix_store_io.rs
View file

@ -17,7 +17,7 @@ use std::{
path::{Path, PathBuf},
sync::Arc,
};
use tokio::io::AsyncReadExt;
use tokio_util::io::SyncIoBridge;
use tracing::{error, instrument, warn, Level};
use tvix_build::buildservice::BuildService;
use tvix_eval::{ErrorKind, EvalIO, FileType, StdIO};
@ -478,7 +478,7 @@ impl EvalIO for TvixStoreIO {
}
#[instrument(skip(self), err)]
fn read_to_end(&self, path: &Path) -> io::Result<Vec<u8>> {
fn open(&self, path: &Path) -> io::Result<Box<dyn io::Read>> {
if let Ok((store_path, sub_path)) =
StorePath::from_absolute_path_full(&path.to_string_lossy())
{
@ -509,10 +509,13 @@ impl EvalIO for TvixStoreIO {
})?;
self.tokio_handle.block_on(async {
let mut reader = {
let resp = self.blob_service.as_ref().open_read(&digest).await?;
match resp {
Some(blob_reader) => blob_reader,
Some(blob_reader) => {
// The VM Response needs a sync [std::io::Reader].
Ok(Box::new(SyncIoBridge::new(blob_reader))
as Box<dyn io::Read>)
}
None => {
error!(
blob.digest = %digest,
@ -521,15 +524,9 @@ impl EvalIO for TvixStoreIO {
Err(io::Error::new(
io::ErrorKind::NotFound,
format!("blob {} not found", &digest),
))?
))
}
}
};
let mut buf = Vec::new();
reader.read_to_end(&mut buf).await?;
Ok(buf)
})
}
Node::Symlink(_symlink_node) => Err(io::Error::new(
@ -540,11 +537,11 @@ impl EvalIO for TvixStoreIO {
} else {
// As tvix-store doesn't manage /nix/store on the filesystem,
// we still need to also ask self.std_io here.
self.std_io.read_to_end(path)
self.std_io.open(path)
}
} else {
// The store path is no store path, so do regular StdIO.
self.std_io.read_to_end(path)
self.std_io.open(path)
}
}