From ed584b92964700b8222d5f97891ed34e319d689c Mon Sep 17 00:00:00 2001 From: Florian Klink Date: Sun, 12 May 2024 13:09:45 +0300 Subject: feat(tvix/castore/directory/traverse_directory): simplify 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 Reviewed-by: Connor Brewster Tested-by: BuildkiteCI --- tvix/castore/src/directoryservice/utils.rs | 91 ++++++++++++++---------------- 1 file changed, 43 insertions(+), 48 deletions(-) (limited to 'tvix/castore/src/directoryservice/utils.rs') diff --git a/tvix/castore/src/directoryservice/utils.rs b/tvix/castore/src/directoryservice/utils.rs index 01c521076c9c..a0ba395ecda8 100644 --- a/tvix/castore/src/directoryservice/utils.rs +++ b/tvix/castore/src/directoryservice/utils.rs @@ -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 = HashSet::new(); - let stream = stream! { + Box::pin(try_stream! { while let Some(current_directory_digest) = worklist_directory_digests.pop_front() { - match directory_service.get(¤t_directory_digest).await { + let current_directory = directory_service.get(¤t_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 + )) - // 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); + // 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 + )) + })?; - // 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 ¤t_directory.directories { - // TODO: propagate error - let child_digest: B3Digest = child_directory_node.digest.clone().try_into().unwrap(); + // 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); - if worklist_directory_digests.contains(&child_digest) - || sent_directory_digests.contains(&child_digest) - { - continue; - } - worklist_directory_digests.push_back(child_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 ¤t_directory.directories { + // TODO: propagate error + let child_digest: B3Digest = child_directory_node.digest.clone().try_into().unwrap(); - yield Ok(current_directory); + if worklist_directory_digests.contains(&child_digest) + || sent_directory_digests.contains(&child_digest) + { + continue; } - }; - } - }; + worklist_directory_digests.push_back(child_digest); + } - Box::pin(stream) + yield current_directory; + } + }) } -- cgit 1.4.1