Add thrift endpoint for get_branches

Reviewed By: StanislavGlebik

Differential Revision: D13765974

fbshipit-source-id: 8c554ca85cde4838673fa6900a6b5643ee9a4973

apiserver/client/src/lib.rs

@@ -14,7 +14,8 @@ use futures_ext::BoxFuture;
 
 use apiserver_thrift::client::{make_MononokeAPIService, MononokeAPIService};
 use apiserver_thrift::types::{
-    MononokeChangeset, MononokeGetChangesetParams, MononokeGetRawParams,
+    MononokeBranches, MononokeChangeset, MononokeGetBranchesParams, MononokeGetChangesetParams,
+    MononokeGetRawParams,
 };
 use srclient::SRChannelBuilder;
 
@@ -55,4 +56,10 @@ impl MononokeAPIClient {
             revision,
         })
     }
+
+    pub fn get_branches(&self) -> BoxFuture<MononokeBranches, apiserver_thrift::errors::Error> {
+        self.inner.get_branches(&MononokeGetBranchesParams {
+            repo: self.repo.clone(),
+        })
+    }
 }

apiserver/client/src/main.rs

@@ -47,6 +47,17 @@ fn get_changeset(client: MononokeAPIClient, matches: &ArgMatches) -> BoxFuture<(
         .boxify()
 }
 
+fn get_branches(client: MononokeAPIClient) -> BoxFuture<(), ()> {
+    client
+        .get_branches()
+        .and_then(|r| {
+            Ok(serde_json::to_string(&r).unwrap_or("Error converting request to json".to_string()))
+        })
+        .map_err(|e| eprintln!("error: {}", e))
+        .map(|res| println!("{}", res))
+        .boxify()
+}
+
 fn main() -> Result<(), ()> {
     let matches = App::new("Mononoke API Server Thrift client")
         .about("Send requests to Mononoke API Server thrift port")
@@ -98,6 +109,7 @@ fn main() -> Result<(), ()> {
                         .required(true),
                 ),
         )
+        .subcommand(SubCommand::with_name("get_branches").about("get all branches"))
         .get_matches();
 
     let tier = matches.value_of("tier").expect("must provide tier name");
@@ -110,6 +122,8 @@ fn main() -> Result<(), ()> {
         cat(client, matches)
     } else if let Some(matches) = matches.subcommand_matches("get_changeset") {
         get_changeset(client, matches)
+    } else if let Some(_) = matches.subcommand_matches("get_branches") {
+        get_branches(client)
     } else {
         Ok(()).into_future().boxify()
     };

apiserver/if/apiserver.thrift

@@ -34,6 +34,13 @@ struct MononokeChangeset {
   6: map<string, binary> extra,
 }
 
+struct MononokeBranches {
+  1: map<string, string> branches,
+}
+
+struct MononokeGetBranchesParams{
+  1: string repo,
+}
 
 service MononokeAPIService extends fb303.FacebookService {
   binary get_raw(1: MononokeGetRawParams params)
@@ -41,4 +48,7 @@ service MononokeAPIService extends fb303.FacebookService {
 
   MononokeChangeset get_changeset(1: MononokeGetChangesetParams param)
     throws (1: MononokeAPIException e),
+
+  MononokeBranches get_branches(1: MononokeGetBranchesParams params)
+    throws (1: MononokeAPIException e),
 }

apiserver/src/actor/query.rs

@@ -11,7 +11,9 @@ use failure::Error;
 
 use http::uri::Uri;
 
-use apiserver_thrift::types::{MononokeGetChangesetParams, MononokeGetRawParams};
+use apiserver_thrift::types::{
+    MononokeGetBranchesParams, MononokeGetChangesetParams, MononokeGetRawParams,
+};
 
 use super::lfs::BatchRequest;
 
@@ -34,6 +36,7 @@ pub enum MononokeRepoQuery {
     GetChangeset {
         revision: String,
     },
+    GetBranches,
     IsAncestor {
         proposed_ancestor: String,
         proposed_descendent: String,
@@ -83,3 +86,14 @@ impl TryFrom<MononokeGetChangesetParams> for MononokeQuery {
         })
     }
 }
+
+impl TryFrom<MononokeGetBranchesParams> for MononokeQuery {
+    type Error = Error;
+
+    fn try_from(params: MononokeGetBranchesParams) -> Result<MononokeQuery, Self::Error> {
+        Ok(MononokeQuery {
+            repo: params.repo,
+            kind: MononokeRepoQuery::GetBranches,
+        })
+    }
+}

apiserver/src/actor/repo.rs

@@ -10,6 +10,7 @@ use bytes::Bytes;
 use failure::{err_msg, Error};
 use futures::future::join_all;
 use futures::sync::oneshot;
+use futures::Stream;
 use futures::{Future, IntoFuture};
 use futures_ext::{BoxFuture, FutureExt};
 use http::uri::Uri;
@@ -268,6 +269,18 @@ impl MononokeRepo {
             .boxify()
     }
 
+    fn get_branches(&self, ctx: CoreContext) -> BoxFuture<MononokeRepoResponse, ErrorKind> {
+        self.repo
+            .get_bookmarks_maybe_stale(ctx)
+            .map(|(bookmark, changesetid)| (bookmark.to_string(), changesetid.to_hex().to_string()))
+            .collect()
+            .map(|vec| MononokeRepoResponse::GetBranches {
+                branches: vec.into_iter().collect(),
+            })
+            .from_err()
+            .boxify()
+    }
+
     fn download_large_file(
         &self,
         ctx: CoreContext,
@@ -342,6 +355,7 @@ impl MononokeRepo {
             ListDirectory { revision, path } => self.list_directory(ctx, revision, path),
             GetTree { hash } => self.get_tree(ctx, hash),
             GetChangeset { revision } => self.get_changeset(ctx, revision),
+            GetBranches => self.get_branches(ctx),
             IsAncestor {
                 proposed_ancestor,
                 proposed_descendent,

apiserver/src/actor/response.rs

@@ -9,6 +9,7 @@ use std::result::Result;
 use actix_web;
 use actix_web::{Body, HttpRequest, HttpResponse, Json, Responder};
 use bytes::Bytes;
+use std::collections::BTreeMap;
 
 use super::lfs::BatchResponse;
 use super::model::{Changeset, Entry, EntryWithSizeAndContentHash};
@@ -29,6 +30,9 @@ pub enum MononokeRepoResponse {
     GetChangeset {
         changeset: Changeset,
     },
+    GetBranches {
+        branches: BTreeMap<String, String>,
+    },
     IsAncestor {
         answer: bool,
     },
@@ -59,6 +63,7 @@ impl Responder for MononokeRepoResponse {
             ListDirectory { files } => Json(files.collect::<Vec<_>>()).respond_to(req),
             GetTree { files } => Json(files).respond_to(req),
             GetChangeset { changeset } => Json(changeset).respond_to(req),
+            GetBranches { branches } => Json(branches).respond_to(req),
             IsAncestor { answer } => Ok(binary_response({
                 if answer {
                     "true".into()

apiserver/src/thrift/mononoke.rs

@@ -8,9 +8,10 @@ use std::convert::TryInto;
 use std::sync::Arc;
 
 use apiserver_thrift::server::MononokeApiservice;
-use apiserver_thrift::services::mononoke_apiservice::{GetChangesetExn, GetRawExn};
+use apiserver_thrift::services::mononoke_apiservice::{GetBranchesExn, GetChangesetExn, GetRawExn};
 use apiserver_thrift::types::{
-    MononokeChangeset, MononokeGetChangesetParams, MononokeGetRawParams,
+    MononokeBranches, MononokeChangeset, MononokeGetBranchesParams, MononokeGetChangesetParams,
+    MononokeGetRawParams,
 };
 use errors::ErrorKind;
 use failure::err_msg;
@@ -18,6 +19,7 @@ use futures::{Future, IntoFuture};
 use futures_ext::BoxFuture;
 use futures_stats::Timed;
 use scuba_ext::{ScubaSampleBuilder, ScubaSampleBuilderExt};
+use serde::Serialize;
 use slog::Logger;
 use time_ext::DurationExt;
 
@@ -46,27 +48,41 @@ impl MononokeAPIServiceImpl {
             scuba_builder,
         }
     }
-}
 
-impl MononokeApiservice for MononokeAPIServiceImpl {
-    fn get_raw(&self, params: MononokeGetRawParams) -> BoxFuture<Vec<u8>, GetRawExn> {
+    fn create_scuba_logger<K: Serialize>(
+        &self,
+        method: &str,
+        params_json: &K,
+        path: Vec<u8>,
+        revision: String,
+    ) -> ScubaSampleBuilder {
         let mut scuba = self.scuba_builder.clone();
-
         scuba
             .add_common_server_data()
             .add("type", "thrift")
-            .add("method", "get_raw")
+            .add("method", method)
             .add(
                 "params",
-                serde_json::to_string(&params)
+                serde_json::to_string(params_json)
                     .unwrap_or_else(|_| "Error converting request to json".to_string()),
             )
             .add(
                 "path",
-                String::from_utf8(params.path.clone())
-                    .unwrap_or("Invalid UTF-8 in path".to_string()),
+                String::from_utf8(path).unwrap_or("Invalid UTF-8 in path".to_string()),
             )
-            .add("revision", params.changeset.clone());
+            .add("revision", revision);
+        scuba
+    }
+}
+
+impl MononokeApiservice for MononokeAPIServiceImpl {
+    fn get_raw(&self, params: MononokeGetRawParams) -> BoxFuture<Vec<u8>, GetRawExn> {
+        let mut scuba = self.create_scuba_logger(
+            "get_raw",
+            &params,
+            params.path.clone(),
+            params.changeset.clone(),
+        );
 
         params
             .try_into()
@@ -101,18 +117,12 @@ impl MononokeApiservice for MononokeAPIServiceImpl {
         &self,
         params: MononokeGetChangesetParams,
     ) -> BoxFuture<MononokeChangeset, GetChangesetExn> {
-        let mut scuba = self.scuba_builder.clone();
-
-        scuba
-            .add_common_server_data()
-            .add("type", "thrift")
-            .add("method", "get_changeset")
-            .add(
-                "params",
-                serde_json::to_string(&params)
-                    .unwrap_or("Error converting request to json".to_string()),
-            )
-            .add("revision", params.revision.clone());
+        let mut scuba = self.create_scuba_logger(
+            "get_changeset",
+            &params,
+            Vec::new(),
+            params.revision.clone(),
+        );
 
         params
             .try_into()
@@ -153,4 +163,49 @@ impl MononokeApiservice for MononokeAPIServiceImpl {
                 }
             })
     }
+
+    fn get_branches(
+        &self,
+        params: MononokeGetBranchesParams,
+    ) -> BoxFuture<MononokeBranches, GetBranchesExn> {
+        let mut scuba =
+            self.create_scuba_logger("get_branches", &params, Vec::new(), String::new());
+
+        params
+            .try_into()
+            .into_future()
+            .from_err()
+            .and_then({
+                cloned!(self.addr);
+                move |param| addr.send_query(param)
+            })
+            .and_then(|resp: MononokeRepoResponse| match resp {
+                MononokeRepoResponse::GetBranches { branches } => Ok(MononokeBranches { branches }),
+                _ => Err(ErrorKind::InternalError(err_msg(
+                    "Actor returned wrong response type to query".to_string(),
+                ))),
+            })
+            .map_err(move |e| GetBranchesExn::e(e.into()))
+            .timed({
+                move |stats, resp| {
+                    scuba
+                        .add(
+                            "response_size",
+                            resp.map(|resp| {
+                                resp.branches
+                                    .iter()
+                                    .map(|(bookmark, hash)| bookmark.len() + hash.len())
+                                    .sum()
+                            })
+                            .unwrap_or(0),
+                        )
+                        .add_future_stats(&stats)
+                        .add("response_time", stats.completion_time.as_micros_unchecked());
+
+                    scuba.log();
+
+                    Ok(())
+                }
+            })
+    }
 }