mononoke/repo_client: add designated manifest fetching in gettreepack

Summary:
This will allow the hg client to do tree fetching like we do in the API Server,
but through the SSH protocol — i.e. by passing a series a manifest ids and
their paths, without recursion on the server side through gettreepack.

Reviewed By: StanislavGlebik

Differential Revision: D20307442

fbshipit-source-id: a6dca03622becdebf41b264381fdd5837a7d4292

eden/mononoke/hgproto/src/lib.rs

@@ -165,7 +165,7 @@ pub struct GettreepackArgs {
     /// "root of the repo".
     pub rootdir: Bytes,
     /// The manifest nodes of the specified root directory to send.
-    pub mfnodes: BTreeSet<HgManifestId>,
+    pub mfnodes: Vec<HgManifestId>,
     /// The manifest nodes of the rootdir that are already on the client.
     pub basemfnodes: BTreeSet<HgManifestId>,
     /// The fullpath (not relative path) of directories underneath

eden/mononoke/hgproto/src/sshproto/request.rs

@@ -598,7 +598,7 @@ fn parse_with_params(
         | call!(parse_command, "gettreepack", parse_params, 0+1,
             |kv| Ok(Gettreepack(GettreepackArgs {
                 rootdir: parseval(&kv, "rootdir", bytes_complete)?,
-                mfnodes: parseval(&kv, "mfnodes", manifestlist)?.into_iter().collect(),
+                mfnodes: parseval(&kv, "mfnodes", manifestlist)?,
                 basemfnodes: parseval(&kv, "basemfnodes", manifestlist)?.into_iter().collect(),
                 directories: parseval(&kv, "directories", gettreepack_directories)?,
                 depth: parseval_option(&kv, "depth", closure!(
@@ -1413,7 +1413,7 @@ mod test_parse {
             inp,
             Request::Single(SingleRequest::Gettreepack(GettreepackArgs {
                 rootdir: Bytes::new(),
-                mfnodes: btreeset![hash_ones_manifest()],
+                mfnodes: vec![hash_ones_manifest()],
                 basemfnodes: btreeset![hash_ones_manifest()],
                 directories: vec![],
                 depth: None,
@@ -1437,7 +1437,7 @@ mod test_parse {
             inp,
             Request::Single(SingleRequest::Gettreepack(GettreepackArgs {
                 rootdir: Bytes::from("ololo".as_bytes()),
-                mfnodes: btreeset![hash_ones_manifest(), hash_twos_manifest()],
+                mfnodes: vec![hash_ones_manifest(), hash_twos_manifest()],
                 basemfnodes: btreeset![hash_twos_manifest(), hash_ones_manifest()],
                 directories: vec![Bytes::from(",".as_bytes()), Bytes::from(";".as_bytes())],
                 depth: Some(1),

eden/mononoke/repo_client/src/client/mod.rs

@@ -180,6 +180,7 @@ fn wireprotocaps() -> Vec<String> {
         "streamreqs=generaldelta,lz4revlog,revlogv1".to_string(),
         "treeonly".to_string(),
         "knownnodes".to_string(),
+        "designatednodes".to_string(),
     ]
 }
 
@@ -1828,19 +1829,57 @@ pub fn gettreepack_entries(
     repo: &BlobRepo,
     params: GettreepackArgs,
 ) -> BoxStream<(HgManifestId, Option<MPath>), Error> {
-    if !params.directories.is_empty() {
-        // This param is not used by core hg, don't worry about implementing it now
-        return stream::once(Err(Error::msg("directories param is not supported"))).boxify();
-    }
-
     let GettreepackArgs {
         rootdir,
         mfnodes,
         basemfnodes,
         depth: fetchdepth,
-        directories: _,
+        directories,
     } = params;
 
+    if fetchdepth == Some(1) && directories.len() > 0 {
+        if directories.len() != mfnodes.len() {
+            let e = format_err!(
+                "invalid directories count ({}, expected {})",
+                directories.len(),
+                mfnodes.len()
+            );
+            return stream::once(Err(e)).boxify();
+        }
+
+        if !rootdir.is_empty() {
+            let e = Error::msg("rootdir must be empty");
+            return stream::once(Err(e)).boxify();
+        }
+
+        if !basemfnodes.is_empty() {
+            let e = Error::msg("basemfnodes must be empty");
+            return stream::once(Err(e)).boxify();
+        }
+
+        let entries = mfnodes
+            .into_iter()
+            .zip(directories.into_iter())
+            .map(|(node, path)| {
+                let path = if path.len() > 0 {
+                    Some(MPath::new(path.as_ref())?)
+                } else {
+                    None
+                };
+                Ok((node, path))
+            })
+            .collect::<Result<Vec<_>, Error>>();
+
+        let entries = try_boxstream!(entries);
+
+        return stream::iter_ok::<_, Error>(entries).boxify();
+    }
+
+    if !directories.is_empty() {
+        // This param is not used by core hg, don't worry about implementing it now
+        return stream::once(Err(Error::msg("directories param is not supported"))).boxify();
+    }
+
     // 65536 matches the default TREE_DEPTH_MAX value from Mercurial
     let fetchdepth = fetchdepth.unwrap_or(2 << 16);