1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
|
use super::Directory;
use super::DirectoryService;
use crate::B3Digest;
use crate::Error;
use crate::Node;
use async_stream::try_stream;
use futures::stream::BoxStream;
use std::collections::{HashSet, VecDeque};
use tracing::instrument;
use tracing::warn;
/// Traverses a [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,
) -> BoxStream<'a, Result<Directory, Error>> {
// The list of all directories that still need to be traversed. The next
// element is picked from the front, new elements are enqueued at the
// back.
let mut worklist_directory_digests: VecDeque<B3Digest> =
VecDeque::from([root_directory_digest.clone()]);
// The list of directory digests already sent to the consumer.
// We omit sending the same directories multiple times.
let mut sent_directory_digests: HashSet<B3Digest> = HashSet::new();
let root_directory_digest = root_directory_digest.clone();
Box::pin(try_stream! {
while let Some(current_directory_digest) = worklist_directory_digests.pop_front() {
let current_directory = match 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
))
})? {
// the root node of the requested closure was not found, return an empty list
None if current_directory_digest == root_directory_digest => break,
// if a child directory of the closure is not there, we have an inconsistent store!
None => {
warn!("directory {} does not exist", current_directory_digest);
Err(Error::StorageError(format!(
"directory {} does not exist",
current_directory_digest
)))?;
break;
}
Some(dir) => dir,
};
// 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.nodes() {
if let Node::Directory{digest: child_digest, ..} = child_directory_node {
if worklist_directory_digests.contains(child_digest)
|| sent_directory_digests.contains(child_digest)
{
continue;
}
worklist_directory_digests.push_back(child_digest.clone());
}
}
yield current_directory;
}
})
}
|
