snix_castore/directoryservice/
utils.rs

1use super::Directory;
2use super::DirectoryService;
3use crate::B3Digest;
4use crate::Error;
5use crate::Node;
6use async_stream::try_stream;
7use futures::stream::BoxStream;
8use std::collections::{HashSet, VecDeque};
9use tracing::instrument;
10use tracing::warn;
11
12/// Traverses a [Directory] from the root to the children.
13///
14/// This is mostly BFS, but directories are only returned once.
15#[instrument(skip(directory_service))]
16pub fn traverse_directory<'a, DS: DirectoryService + 'static>(
17    directory_service: DS,
18    root_directory_digest: &B3Digest,
19) -> BoxStream<'a, Result<Directory, Error>> {
20    // The list of all directories that still need to be traversed. The next
21    // element is picked from the front, new elements are enqueued at the
22    // back.
23    let mut worklist_directory_digests: VecDeque<B3Digest> =
24        VecDeque::from([root_directory_digest.clone()]);
25    // The list of directory digests already sent to the consumer.
26    // We omit sending the same directories multiple times.
27    let mut sent_directory_digests: HashSet<B3Digest> = HashSet::new();
28
29    let root_directory_digest = root_directory_digest.clone();
30
31    Box::pin(try_stream! {
32        while let Some(current_directory_digest) = worklist_directory_digests.pop_front() {
33            let current_directory = match directory_service.get(&current_directory_digest).await.map_err(|e| {
34                warn!("failed to look up directory");
35                Error::StorageError(format!(
36                    "unable to look up directory {}: {}",
37                    current_directory_digest, e
38                ))
39            })? {
40                // the root node of the requested closure was not found, return an empty list
41                None if current_directory_digest == root_directory_digest => break,
42                // if a child directory of the closure is not there, we have an inconsistent store!
43                None => {
44                    warn!("directory {} does not exist", current_directory_digest);
45                    Err(Error::StorageError(format!(
46                        "directory {} does not exist",
47                        current_directory_digest
48                    )))?;
49                    break;
50                }
51                Some(dir) => dir,
52            };
53
54            // We're about to send this directory, so let's avoid sending it again if a
55            // descendant has it.
56            sent_directory_digests.insert(current_directory_digest);
57
58            // enqueue all child directory digests to the work queue, as
59            // long as they're not part of the worklist or already sent.
60            // This panics if the digest looks invalid, it's supposed to be checked first.
61            for (_, child_directory_node) in current_directory.nodes() {
62                if let Node::Directory{digest: child_digest, ..} = child_directory_node {
63                    if worklist_directory_digests.contains(child_digest)
64                        || sent_directory_digests.contains(child_digest)
65                    {
66                        continue;
67                    }
68                    worklist_directory_digests.push_back(child_digest.clone());
69                }
70            }
71
72            yield current_directory;
73        }
74    })
75}