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

feat(tvix/castore/directory/traverse_directory): simplify

Browse source 
Use try_stream! rather than stream!, and a bit more map_err and ok_err
to make things a bit more concise. Once we have proper error types here,
and impl Froms, a lot of the error mapping would disappear entirely.

Change-Id: I5240a6b0ff7818b94c151322774242b2c142e33b
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11633
Autosubmit: flokli <flokli@flokli.de>
Reviewed-by: Connor Brewster <cbrewster@hey.com>
Tested-by: BuildkiteCI
This commit is contained in:
Florian Klink 2024-05-12 13:09:45 +03:00 committed by clbot
parent 7fd4adc129
commit ed584b9296
1 changed files with 47 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

99
tvix/castore/src/directoryservice/utils.rs
View file

@ -2,14 +2,16 @@ use super::DirectoryService;
use crate::proto;
use crate::B3Digest;
use crate::Error;
use async_stream::stream;
use async_stream::try_stream;
use futures::stream::BoxStream;
use std::collections::{HashSet, VecDeque};
use tracing::instrument;
use tracing::warn;
/// Traverses a [proto::Directory] from the root to the children.
///
/// This is mostly BFS, but directories are only returned once.
#[instrument(skip(directory_service))]
pub fn traverse_directory<'a, DS: DirectoryService + 'static>(
directory_service: DS,
root_directory_digest: &B3Digest,
@ -23,60 +25,53 @@ pub fn traverse_directory<'a, DS: DirectoryService + 'static>(
// We omit sending the same directories multiple times.
let mut sent_directory_digests: HashSet<B3Digest> = HashSet::new();
let stream = stream! {
Box::pin(try_stream! {
while let Some(current_directory_digest) = worklist_directory_digests.pop_front() {
match directory_service.get(&current_directory_digest).await {
let current_directory = directory_service.get(&current_directory_digest).await.map_err(|e| {
warn!("failed to look up directory");
Error::StorageError(format!(
"unable to look up directory {}: {}",
current_directory_digest, e
))
})?.ok_or_else(|| {
// if it's not there, we have an inconsistent store!
Ok(None) => {
warn!("directory {} does not exist", current_directory_digest);
yield Err(Error::StorageError(format!(
"directory {} does not exist",
current_directory_digest
)));
}
Err(e) => {
warn!("failed to look up directory");
yield Err(Error::StorageError(format!(
"unable to look up directory {}: {}",
current_directory_digest, e
)));
warn!("directory {} does not exist", current_directory_digest);
Error::StorageError(format!(
"directory {} does not exist",
current_directory_digest
))
})?;
// validate, we don't want to send invalid directories.
current_directory.validate().map_err(|e| {
warn!("directory failed validation: {}", e.to_string());
Error::StorageError(format!(
"invalid directory: {}",
current_directory_digest
))
})?;
// We're about to send this directory, so let's avoid sending it again if a
// descendant has it.
sent_directory_digests.insert(current_directory_digest);
// enqueue all child directory digests to the work queue, as
// long as they're not part of the worklist or already sent.
// This panics if the digest looks invalid, it's supposed to be checked first.
for child_directory_node in &current_directory.directories {
// TODO: propagate error
let child_digest: B3Digest = child_directory_node.digest.clone().try_into().unwrap();
if worklist_directory_digests.contains(&child_digest)
|| sent_directory_digests.contains(&child_digest)
{
continue;
}
worklist_directory_digests.push_back(child_digest);
}
// if we got it
Ok(Some(current_directory)) => {
// validate, we don't want to send invalid directories.
if let Err(e) = current_directory.validate() {
warn!("directory failed validation: {}", e.to_string());
yield Err(Error::StorageError(format!(
"invalid directory: {}",
current_directory_digest
)));
}
// We're about to send this directory, so let's avoid sending it again if a
// descendant has it.
sent_directory_digests.insert(current_directory_digest);
// enqueue all child directory digests to the work queue, as
// long as they're not part of the worklist or already sent.
// This panics if the digest looks invalid, it's supposed to be checked first.
for child_directory_node in &current_directory.directories {
// TODO: propagate error
let child_digest: B3Digest = child_directory_node.digest.clone().try_into().unwrap();
if worklist_directory_digests.contains(&child_digest)
|| sent_directory_digests.contains(&child_digest)
{
continue;
}
worklist_directory_digests.push_back(child_digest);
}
yield Ok(current_directory);
}
};
yield current_directory;
}
};
Box::pin(stream)
})
}