refactor(tvix/store): Asyncify PathInfoService and DirectoryService r/6623

We've decided to asyncify all of the services to reduce some of the
pains going back and for between sync<->async. The end goal will be for
all the tvix-store internals to be async and then expose a sync
interface for things like tvix eval io.

Change-Id: I97c71f8db1d05a38bd8f625df5087d565705d52d
Reviewed-on: https://cl.tvl.fyi/c/depot/+/9369
Autosubmit: Connor Brewster <cbrewster@hey.com>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

1 files changed, 64 insertions, 51 deletions

diff --git a/tvix/store/src/directoryservice/traverse.rs b/tvix/store/src/directoryservice/traverse.rs
index 9029543267c1..0489c4458163 100644
--- a/tvix/store/src/directoryservice/traverse.rs
+++ b/tvix/store/src/directoryservice/traverse.rs
@@ -1,19 +1,18 @@
 use super::DirectoryService;
-use crate::{proto::NamedNode, Error};
+use crate::{proto::NamedNode, B3Digest, Error};
 use std::{os::unix::ffi::OsStrExt, sync::Arc};
 use tracing::{instrument, warn};
 
 /// This traverses from a (root) node to the given (sub)path, returning the Node
 /// at that path, or none, if there's nothing at that path.
-/// TODO: Do we want to rewrite this in a non-recursing fashion, and use
-/// [DirectoryService.get_recursive] to do less lookups?
+/// TODO: Do we want to use [DirectoryService.get_recursive] to do less lookups?
 /// Or do we consider this to be a non-issue due to store composition and local caching?
 /// TODO: the name of this function (and mod) is a bit bad, because it doesn't
 /// clearly distinguish it from the BFS traversers.
 #[instrument(skip(directory_service))]
-pub fn traverse_to(
+pub async fn traverse_to(
     directory_service: Arc<dyn DirectoryService>,
-    node: crate::proto::node::Node,
+    root_node: crate::proto::node::Node,
     path: &std::path::Path,
 ) -> Result<Option<crate::proto::node::Node>, Error> {
     // strip a possible `/` prefix from the path.
@@ -25,52 +24,56 @@ pub fn traverse_to(
         }
     };
 
+    let mut cur_node = root_node;
     let mut it = path.components();
 
-    match it.next() {
-        None => {
-            // the (remaining) path is empty, return the node we've been called with.
-            Ok(Some(node))
-        }
-        Some(first_component) => {
-            match node {
-                crate::proto::node::Node::File(_) | crate::proto::node::Node::Symlink(_) => {
-                    // There's still some path left, but the current node is no directory.
-                    // This means the path doesn't exist, as we can't reach it.
-                    Ok(None)
-                }
-                crate::proto::node::Node::Directory(directory_node) => {
-                    let digest = directory_node
-                        .digest
-                        .try_into()
-                        .map_err(|_e| Error::StorageError("invalid digest length".to_string()))?;
-
-                    // fetch the linked node from the directory_service
-                    match directory_service.get(&digest)? {
-                        // If we didn't get the directory node that's linked, that's a store inconsistency, bail out!
-                        None => {
-                            warn!("directory {} does not exist", digest);
-
-                            Err(Error::StorageError(format!(
-                                "directory {} does not exist",
-                                digest
-                            )))
-                        }
-                        Some(directory) => {
-                            // look for first_component in the [Directory].
-                            // FUTUREWORK: as the nodes() iterator returns in a sorted fashion, we
-                            // could stop as soon as e.name is larger than the search string.
-                            let child_node = directory
-                                .nodes()
-                                .find(|n| n.get_name() == first_component.as_os_str().as_bytes());
-
-                            match child_node {
-                                // child node not found means there's no such element inside the directory.
-                                None => Ok(None),
-                                // child node found, recurse with it and the rest of the path.
-                                Some(child_node) => {
-                                    let rest_path: std::path::PathBuf = it.collect();
-                                    traverse_to(directory_service, child_node, &rest_path)
+    loop {
+        match it.next() {
+            None => {
+                // the (remaining) path is empty, return the node we're current at.
+                return Ok(Some(cur_node));
+            }
+            Some(first_component) => {
+                match cur_node {
+                    crate::proto::node::Node::File(_) | crate::proto::node::Node::Symlink(_) => {
+                        // There's still some path left, but the current node is no directory.
+                        // This means the path doesn't exist, as we can't reach it.
+                        return Ok(None);
+                    }
+                    crate::proto::node::Node::Directory(directory_node) => {
+                        let digest: B3Digest = directory_node.digest.try_into().map_err(|_e| {
+                            Error::StorageError("invalid digest length".to_string())
+                        })?;
+
+                        // fetch the linked node from the directory_service
+                        match directory_service.get(&digest).await? {
+                            // If we didn't get the directory node that's linked, that's a store inconsistency, bail out!
+                            None => {
+                                warn!("directory {} does not exist", digest);
+
+                                return Err(Error::StorageError(format!(
+                                    "directory {} does not exist",
+                                    digest
+                                )));
+                            }
+                            Some(directory) => {
+                                // look for first_component in the [Directory].
+                                // FUTUREWORK: as the nodes() iterator returns in a sorted fashion, we
+                                // could stop as soon as e.name is larger than the search string.
+                                let child_node = directory.nodes().find(|n| {
+                                    n.get_name() == first_component.as_os_str().as_bytes()
+                                });
+
+                                match child_node {
+                                    // child node not found means there's no such element inside the directory.
+                                    None => {
+                                        return Ok(None);
+                                    }
+                                    // child node found, return to top-of loop to find the next
+                                    // node in the path.
+                                    Some(child_node) => {
+                                        cur_node = child_node;
+                                    }
                                 }
                             }
                         }
@@ -92,16 +95,18 @@ mod tests {
 
     use super::traverse_to;
 
-    #[test]
-    fn test_traverse_to() {
+    #[tokio::test]
+    async fn test_traverse_to() {
         let directory_service = gen_directory_service();
 
         let mut handle = directory_service.put_multiple_start();
         handle
             .put(DIRECTORY_WITH_KEEP.clone())
+            .await
             .expect("must succeed");
         handle
             .put(DIRECTORY_COMPLICATED.clone())
+            .await
             .expect("must succeed");
 
         // construct the node for DIRECTORY_COMPLICATED
@@ -128,6 +133,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from(""),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(Some(node_directory_complicated.clone()), resp);
@@ -140,6 +146,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from("keep"),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(Some(node_directory_with_keep), resp);
@@ -152,6 +159,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from("keep/.keep"),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(Some(node_file_keep.clone()), resp);
@@ -164,6 +172,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from("/keep/.keep"),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(Some(node_file_keep), resp);
@@ -176,6 +185,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from("void"),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(None, resp);
@@ -188,6 +198,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from("//v/oid"),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(None, resp);
@@ -201,6 +212,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from("keep/.keep/foo"),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(None, resp);
@@ -213,6 +225,7 @@ mod tests {
                 node_directory_complicated.clone(),
                 &PathBuf::from("/"),
             )
+            .await
             .expect("must succeed");
 
             assert_eq!(Some(node_directory_complicated), resp);