Add tree prefetch endpoint

Summary: Add an endpoint to the API server that provides functionality similar to the `gettreepack` wire protocol command.

Reviewed By: fanzeyi

Differential Revision: D15492734

fbshipit-source-id: 7d0f113f0d33c68d5bfba5781269a92f0d5a66e8

apiserver/src/actor/query.rs

@@ -17,7 +17,7 @@ use apiserver_thrift::types::{
     MononokeGetRawParams, MononokeGetTreeParams, MononokeIsAncestorParams,
     MononokeListDirectoryParams, MononokeRevision,
 };
-use types::api::{DataRequest, HistoryRequest};
+use types::api::{DataRequest, HistoryRequest, TreeRequest};
 
 use super::lfs::BatchRequest;
 
@@ -74,6 +74,7 @@ pub enum MononokeRepoQuery {
     EdenGetData(DataRequest),
     EdenGetHistory(HistoryRequest),
     EdenGetTrees(DataRequest),
+    EdenPrefetchTrees(TreeRequest),
 }
 
 pub struct MononokeQuery {

apiserver/src/actor/repo.rs

@@ -29,6 +29,7 @@ use futures_ext::{try_boxfuture, BoxFuture, FutureExt, StreamExt};
 use http::uri::Uri;
 use mononoke_api;
 use remotefilelog;
+use repo_client::gettreepack_entries;
 use scuba_ext::ScubaSampleBuilder;
 use slog::{debug, Logger};
 use sshrelay::SshEnvVars;
@@ -38,8 +39,8 @@ use uuid::Uuid;
 use mercurial_types::{manifest::Content, Entry as _, HgChangesetId, HgFileNodeId, HgManifestId};
 use metaconfig_types::RepoConfig;
 use types::{
-    api::{DataRequest, DataResponse, HistoryRequest, HistoryResponse},
-    DataEntry, Key, WireHistoryEntry,
+    api::{DataRequest, DataResponse, HistoryRequest, HistoryResponse, TreeRequest},
+    DataEntry, Key, RepoPathBuf, WireHistoryEntry,
 };
 
 use mononoke_types::{FileContents, MPath, RepositoryId};
@@ -533,6 +534,38 @@ impl MononokeRepo {
             .boxify()
     }
 
+    fn eden_prefetch_trees(
+        &self,
+        ctx: CoreContext,
+        req: TreeRequest,
+    ) -> BoxFuture<MononokeRepoResponse, ErrorKind> {
+        gettreepack_entries(ctx.clone(), &self.repo, req.into())
+            .and_then(move |(entry, basepath)| {
+                let full_path = MPath::join_element_opt(basepath.as_ref(), entry.get_name());
+                let path_bytes = full_path
+                    .map(|mpath| mpath.to_vec())
+                    .unwrap_or_else(Vec::new);
+                let path = try_boxfuture!(RepoPathBuf::from_utf8(path_bytes));
+
+                let node = entry.get_hash().into_nodehash().into();
+                let key = Key::new(path, node);
+
+                let get_parents = entry.get_parents(ctx.clone());
+                let get_content = entry.get_raw_content(ctx.clone());
+                get_parents
+                    .and_then(move |parents| {
+                        get_content.map(move |content| {
+                            DataEntry::new(key, content.into_inner(), parents.into())
+                        })
+                    })
+                    .boxify()
+            })
+            .collect()
+            .map(|entries| MononokeRepoResponse::EdenPrefetchTrees(DataResponse::new(entries)))
+            .from_err()
+            .boxify()
+    }
+
     pub fn send_query(
         &self,
         ctx: CoreContext,
@@ -570,6 +603,7 @@ impl MononokeRepo {
                 self.eden_get_history(ctx, keys, depth)
             }
             EdenGetTrees(DataRequest { keys }) => self.eden_get_trees(ctx, keys),
+            EdenPrefetchTrees(req) => self.eden_prefetch_trees(ctx, req),
         }
     }
 }

apiserver/src/actor/response.rs

@@ -57,6 +57,7 @@ pub enum MononokeRepoResponse {
     EdenGetData(DataResponse),
     EdenGetHistory(HistoryResponse),
     EdenGetTrees(DataResponse),
+    EdenPrefetchTrees(DataResponse),
 }
 
 fn binary_response(content: Bytes) -> HttpResponse {
@@ -107,6 +108,7 @@ impl Responder for MononokeRepoResponse {
             EdenGetData(response) => Ok(cbor_response(response)),
             EdenGetHistory(response) => Ok(cbor_response(response)),
             EdenGetTrees(response) => Ok(cbor_response(response)),
+            EdenPrefetchTrees(response) => Ok(cbor_response(response)),
         }
     }
 }

apiserver/src/main.rs

@@ -424,6 +424,34 @@ fn eden_get_trees(
     }
 }
 
+#[derive(Deserialize)]
+struct EdenPrefetchTreesParams {
+    repo: String,
+}
+
+fn eden_prefetch_trees(
+    (state, params, body): (State<HttpServerState>, Path<EdenPrefetchTreesParams>, Bytes),
+) -> impl Future<Item = MononokeRepoResponse, Error = ErrorKind> {
+    let params = params.into_inner();
+    match serde_cbor::from_slice(&body) {
+        Ok(request) => state
+            .mononoke
+            .send_query(
+                prepare_fake_ctx(&state),
+                MononokeQuery {
+                    repo: params.repo,
+                    kind: MononokeRepoQuery::EdenPrefetchTrees(request),
+                },
+            )
+            .left_future(),
+        Err(e) => {
+            let msg = "POST data is invalid CBOR".into();
+            let e = ErrorKind::InvalidInput(msg, Some(e.into()));
+            err(e).right_future()
+        }
+    }
+}
+
 fn setup_logger(debug: bool) -> Logger {
     let level = if debug { Level::Debug } else { Level::Info };
 
@@ -700,6 +728,12 @@ fn main() -> Fallible<()> {
                             (cfg.0).2.limit(config::MAX_PAYLOAD_SIZE);
                         })
                 })
+                .resource("/eden/trees/prefetch", |r| {
+                    r.method(http::Method::POST)
+                        .with_async_config(eden_prefetch_trees, |cfg| {
+                            (cfg.0).2.limit(config::MAX_PAYLOAD_SIZE);
+                        })
+                })
             })
     });
 

repo_client/src/client/mod.rs

@@ -1386,7 +1386,7 @@ impl HgCommands for RepoClient {
     }
 }
 
-fn gettreepack_entries(
+pub fn gettreepack_entries(
     ctx: CoreContext,
     repo: &BlobRepo,
     params: GettreepackArgs,

repo_client/src/lib.rs

@@ -14,7 +14,7 @@ mod client;
 mod errors;
 mod mononoke_repo;
 
-pub use client::RepoClient;
+pub use client::{gettreepack_entries, RepoClient};
 pub use mononoke_repo::{streaming_clone, MononokeRepo};
 pub use repo_read_write_status::RepoReadWriteFetcher;
 pub use streaming_clone::SqlStreamingChunksFetcher;