manifest: add BFS prefetch function

Summary:
Add a client-driven tree prefetching implementation to the Rust manifest code. Unlike the existing prefetch implementation in Python, this one does all computation of which nodes to fetch on the client side using the BFS logic from BfsDiff. The trees are then bulk fetched layer-by-layer using EdenAPI.

This initial version is fairly naive, and omits some obvious optimizations (such as performing fetches of multiple trees concurrently), but is sufficient to demonstrate HTTP tree prefetching in action.

Reviewed By: xavierd

Differential Revision: D17379178

fbshipit-source-id: f17fe99834ad4fec07b4a4ab196928cc4fe91142

edenscm/hgext/treemanifest/__init__.py

@@ -551,6 +551,9 @@ def wraprepo(repo):
             if depth is None:
                 depth = self.ui.configint("treemanifest", "fetchdepth")
 
+            if self._httpprefetchtrees(mfnodes, depth):
+                return
+
             start = util.timer()
             with self.ui.timesection("fetchingtrees"):
                 with self.connectionpool.get(fallbackpath) as conn:
@@ -572,6 +575,20 @@ def wraprepo(repo):
                 caps = _addservercaps(self, caps)
             return caps
 
+        def _httpprefetchtrees(self, mfnodes, depth=None):
+            mfl = self.manifestlog
+            if _userustmanifest(mfl) and _usehttp(self.ui):
+                try:
+                    with progress.spinner(self.ui, "Prefetching trees over HTTPS"):
+                        for node in mfnodes:
+                            rustmanifest.prefetch(mfl.datastore, node, "", depth)
+                    return True
+                except Exception as e:
+                    self.ui.warn(_("encountered error during HTTPS fetching;"))
+                    self.ui.warn(_(" falling back to SSH\n"))
+                    edenapi.logexception(self.ui, e)
+            return False
+
         def _httpgettrees(self, keys):
             """
             Fetch the specified tree nodes over HTTP via the Eden API.
@@ -1020,6 +1037,10 @@ class hybridmanifestlog(manifest.manifestlog):
         self.treemanifestlog.abortpending()
 
 
+def _usehttp(ui):
+    return edenapi.enabled(ui) and ui.configbool("treemanifest", "usehttp")
+
+
 def _userustmanifest(manifestlog):
     return manifestlog.ui.configbool("treemanifest", "rustmanifest")
 
@@ -2529,8 +2550,7 @@ class remotetreestore(generatingdatastore):
         # back to the old prefetch-based fetching behavior if the
         # HTTP fetch fails. Unlike _prefetch(), the HTTP fetching
         # only fetches the individually requested tree node.
-        usehttp = self._repo.ui.configbool("treemanifest", "usehttp")
-        if edenapi.enabled(self._repo.ui) and usehttp:
+        if _usehttp(self._repo.ui):
             try:
                 self._repo._httpgettrees([(name, node)])
                 return
@@ -2578,8 +2598,7 @@ class remotetreestore(generatingdatastore):
         self._sharedhistory.markforrefresh()
 
     def prefetch(self, keys):
-        usehttp = self._repo.ui.configbool("treemanifest", "usehttp")
-        if edenapi.enabled(self._repo.ui) and usehttp:
+        if _usehttp(self._repo.ui):
             keys = self._shareddata.getmissing(keys)
             try:
                 self._repo._httpgettrees(keys)

edenscmnative/bindings/modules/pymanifest/src/lib.rs

@@ -66,6 +66,19 @@ pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
             )
         ),
     )?;
+    m.add(
+        py,
+        "prefetch",
+        py_fn!(
+            py,
+            prefetch(
+                store: PyObject,
+                node: &PyBytes,
+                path: &PyBytes,
+                depth: Option<usize> = None
+            )
+        ),
+    )?;
     Ok(m)
 }
 
@@ -442,6 +455,21 @@ pub fn subdir_diff(
     vec_to_iter(py, result)
 }
 
+pub fn prefetch(
+    py: Python,
+    store: PyObject,
+    node: &PyBytes,
+    path: &PyBytes,
+    depth: Option<usize>,
+) -> PyResult<PyObject> {
+    let store = Arc::new(ManifestStore::new(PythonDataStore::new(store)));
+    let node = pybytes_to_node(py, node)?;
+    let path = pybytes_to_path(py, path);
+    let key = Key::new(path, node);
+    manifest::prefetch(store, key, depth).map_pyerr::<exc::RuntimeError>(py)?;
+    Ok(py.None())
+}
+
 fn vec_to_iter<T: ToPyObject>(py: Python, items: Vec<T>) -> PyResult<PyObject> {
     let list: PyList = items.into_py_object(py);
     list.into_object().call_method(py, "__iter__", NoArgs, None)

lib/manifest/src/lib.rs

@@ -68,4 +68,4 @@ pub enum FsNode {
 mod file;
 pub mod tree;
 pub use crate::file::{FileMetadata, FileType};
-pub use crate::tree::{compat_subtree_diff, Diff, DiffEntry, DiffType, Tree, TreeStore};
+pub use crate::tree::{compat_subtree_diff, prefetch, Diff, DiffEntry, DiffType, Tree, TreeStore};

lib/manifest/src/tree/mod.rs

@@ -568,6 +568,38 @@ pub fn compat_subtree_diff(
     Ok(state.result)
 }
 
+pub fn prefetch(
+    store: Arc<dyn TreeStore + Send + Sync>,
+    key: Key,
+    mut depth: Option<usize>,
+) -> Fallible<()> {
+    let tree = Tree::durable(store, key.node);
+    let mut dirs = vec![Directory::from_link(&tree.root, key.path).unwrap()];
+
+    while !dirs.is_empty() {
+        let keys = dirs.iter().filter_map(|d| d.key()).collect::<Vec<_>>();
+        if !keys.is_empty() {
+            tree.store.prefetch(keys)?;
+        }
+
+        dirs = dirs
+            .into_iter()
+            .map(|d| Ok(d.list(&tree.store)?.1))
+            .collect::<Fallible<Vec<_>>>()?
+            .into_iter()
+            .flatten()
+            .collect();
+
+        depth = match depth {
+            Some(0) => break,
+            Some(d) => Some(d - 1),
+            None => None,
+        };
+    }
+
+    Ok(())
+}
+
 /// A file (leaf node) encountered during a tree traversal.
 ///
 /// Consists of the full path to the file along with the associated file metadata.