diff options
Diffstat (limited to 'tvix/castore/src/directoryservice/directory_graph.rs')
| -rw-r--r-- | tvix/castore/src/directoryservice/directory_graph.rs | 119 |
1 files changed, 55 insertions, 64 deletions
diff --git a/tvix/castore/src/directoryservice/directory_graph.rs b/tvix/castore/src/directoryservice/directory_graph.rs index e6b9b163370c..54f3cb3e9353 100644 --- a/tvix/castore/src/directoryservice/directory_graph.rs +++ b/tvix/castore/src/directoryservice/directory_graph.rs @@ -1,7 +1,5 @@ use std::collections::HashMap; -use bstr::ByteSlice; - use petgraph::{ graph::{DiGraph, NodeIndex}, visit::{Bfs, DfsPostOrder, EdgeRef, IntoNodeIdentifiers, Walker}, @@ -10,10 +8,7 @@ use petgraph::{ use tracing::instrument; use super::order_validator::{LeavesToRootValidator, OrderValidator, RootToLeavesValidator}; -use crate::{ - proto::{self, Directory, DirectoryNode}, - B3Digest, -}; +use crate::{path::PathComponent, B3Digest, Directory, Node}; #[derive(thiserror::Error, Debug)] pub enum Error { @@ -21,6 +16,11 @@ pub enum Error { ValidationError(String), } +struct EdgeWeight { + name: PathComponent, + size: u64, +} + /// This can be used to validate and/or re-order a Directory closure (DAG of /// connected Directories), and their insertion order. /// @@ -58,7 +58,7 @@ pub struct DirectoryGraph<O> { // // The option in the edge weight tracks the pending validation state of the respective edge, for example if // the child has not been added yet. - graph: DiGraph<Option<Directory>, Option<DirectoryNode>>, + graph: DiGraph<Option<Directory>, Option<EdgeWeight>>, // A lookup table from directory digest to node index. digest_to_node_ix: HashMap<B3Digest, NodeIndex>, @@ -67,18 +67,18 @@ pub struct DirectoryGraph<O> { } pub struct ValidatedDirectoryGraph { - graph: DiGraph<Option<Directory>, Option<DirectoryNode>>, + graph: DiGraph<Option<Directory>, Option<EdgeWeight>>, root: Option<NodeIndex>, } -fn check_edge(dir: &DirectoryNode, child: &Directory) -> Result<(), Error> { +fn check_edge(edge: &EdgeWeight, child: &Directory) -> Result<(), Error> { // Ensure the size specified in the child node matches our records. - if dir.size != child.size() { + if edge.size != child.size() { return Err(Error::ValidationError(format!( "'{}' has wrong size, specified {}, recorded {}", - dir.name.as_bstr(), - dir.size, + edge.name, + edge.size, child.size(), ))); } @@ -88,7 +88,7 @@ fn check_edge(dir: &DirectoryNode, child: &Directory) -> Result<(), Error> { impl DirectoryGraph<LeavesToRootValidator> { /// Insert a new Directory into the closure #[instrument(level = "trace", skip_all, fields(directory.digest=%directory.digest(), directory.size=%directory.size()), err)] - pub fn add(&mut self, directory: proto::Directory) -> Result<(), Error> { + pub fn add(&mut self, directory: Directory) -> Result<(), Error> { if !self.order_validator.add_directory(&directory) { return Err(Error::ValidationError( "unknown directory was referenced".into(), @@ -108,7 +108,7 @@ impl DirectoryGraph<RootToLeavesValidator> { /// Insert a new Directory into the closure #[instrument(level = "trace", skip_all, fields(directory.digest=%directory.digest(), directory.size=%directory.size()), err)] - pub fn add(&mut self, directory: proto::Directory) -> Result<(), Error> { + pub fn add(&mut self, directory: Directory) -> Result<(), Error> { let digest = directory.digest(); if !self.order_validator.digest_allowed(&digest) { return Err(Error::ValidationError("unexpected digest".into())); @@ -129,12 +129,7 @@ impl<O: OrderValidator> DirectoryGraph<O> { } /// Adds a directory which has already been confirmed to be in-order to the graph - pub fn add_order_unchecked(&mut self, directory: proto::Directory) -> Result<(), Error> { - // Do some basic validation - directory - .validate() - .map_err(|e| Error::ValidationError(e.to_string()))?; - + pub fn add_order_unchecked(&mut self, directory: Directory) -> Result<(), Error> { let digest = directory.digest(); // Teach the graph about the existence of a node with this digest @@ -149,23 +144,32 @@ impl<O: OrderValidator> DirectoryGraph<O> { } // set up edges to all child directories - for subdir in &directory.directories { - let subdir_digest: B3Digest = subdir.digest.clone().try_into().unwrap(); - - let child_ix = *self - .digest_to_node_ix - .entry(subdir_digest) - .or_insert_with(|| self.graph.add_node(None)); - - let pending_edge_check = match &self.graph[child_ix] { - Some(child) => { - // child is already available, validate the edge now - check_edge(subdir, child)?; - None - } - None => Some(subdir.clone()), // pending validation - }; - self.graph.add_edge(ix, child_ix, pending_edge_check); + for (name, node) in directory.nodes() { + if let Node::Directory { digest, size } = node { + let child_ix = *self + .digest_to_node_ix + .entry(digest.clone()) + .or_insert_with(|| self.graph.add_node(None)); + + let pending_edge_check = match &self.graph[child_ix] { + Some(child) => { + // child is already available, validate the edge now + check_edge( + &EdgeWeight { + name: name.clone(), + size: *size, + }, + child, + )?; + None + } + None => Some(EdgeWeight { + name: name.clone(), + size: *size, + }), // pending validation + }; + self.graph.add_edge(ix, child_ix, pending_edge_check); + } } // validate the edges from parents to this node @@ -183,6 +187,7 @@ impl<O: OrderValidator> DirectoryGraph<O> { .expect("edge not found") .take() .expect("edge is already validated"); + check_edge(&edge_weight, &directory)?; } @@ -269,33 +274,23 @@ impl ValidatedDirectoryGraph { #[cfg(test)] mod tests { - use crate::{ - fixtures::{DIRECTORY_A, DIRECTORY_B, DIRECTORY_C}, - proto::{self, Directory}, - }; - use lazy_static::lazy_static; + use crate::fixtures::{DIRECTORY_A, DIRECTORY_B, DIRECTORY_C}; + use crate::{Directory, Node}; use rstest::rstest; + use std::sync::LazyLock; - lazy_static! { - pub static ref BROKEN_DIRECTORY : Directory = Directory { - symlinks: vec![proto::SymlinkNode { - name: "".into(), // invalid name! - target: "doesntmatter".into(), - }], - ..Default::default() - }; + use super::{DirectoryGraph, LeavesToRootValidator, RootToLeavesValidator}; - pub static ref BROKEN_PARENT_DIRECTORY: Directory = Directory { - directories: vec![proto::DirectoryNode { - name: "foo".into(), - digest: DIRECTORY_A.digest().into(), + pub static BROKEN_PARENT_DIRECTORY: LazyLock<Directory> = LazyLock::new(|| { + Directory::try_from_iter([( + "foo".try_into().unwrap(), + Node::Directory { + digest: DIRECTORY_A.digest(), size: DIRECTORY_A.size() + 42, // wrong! - }], - ..Default::default() - }; - } - - use super::{DirectoryGraph, LeavesToRootValidator, RootToLeavesValidator}; + }, + )]) + .unwrap() + }); #[rstest] /// Uploading an empty directory should succeed. @@ -312,8 +307,6 @@ mod tests { #[case::unconnected_node(&[&*DIRECTORY_A, &*DIRECTORY_C, &*DIRECTORY_B], false, None)] /// Uploading B (referring to A) should fail immediately, because A was never uploaded. #[case::dangling_pointer(&[&*DIRECTORY_B], true, None)] - /// Uploading a directory failing validation should fail immediately. - #[case::failing_validation(&[&*BROKEN_DIRECTORY], true, None)] /// Uploading a directory which refers to another Directory with a wrong size should fail. #[case::wrong_size_in_parent(&[&*DIRECTORY_A, &*BROKEN_PARENT_DIRECTORY], true, None)] fn test_uploads( @@ -366,8 +359,6 @@ mod tests { #[case::unconnected_node(&*DIRECTORY_C, &[&*DIRECTORY_C, &*DIRECTORY_B], true, None)] /// Downloading B (specified as the root) but receiving A instead should fail immediately, because A has no connection to B (the root). #[case::dangling_pointer(&*DIRECTORY_B, &[&*DIRECTORY_A], true, None)] - /// Downloading a directory failing validation should fail immediately. - #[case::failing_validation(&*BROKEN_DIRECTORY, &[&*BROKEN_DIRECTORY], true, None)] /// Downloading a directory which refers to another Directory with a wrong size should fail. #[case::wrong_size_in_parent(&*BROKEN_PARENT_DIRECTORY, &[&*BROKEN_PARENT_DIRECTORY, &*DIRECTORY_A], true, None)] fn test_downloads( |