edenapi: add commit/graph_segments handler

Summary: Adds a new EdenAPI endpoint commit/graph_segments that returns a segmented representation of ancestors of one set of commits (heads) excluding ancestors of another set ot commits (common).

Reviewed By: markbt

Differential Revision: D46796194

fbshipit-source-id: 45513e97a1b42f3fbdbe14ed94961512011cb0d5

eden/mononoke/edenapi_service/src/handlers/commit.rs

@@ -23,6 +23,9 @@ use edenapi_types::Batch;
 use edenapi_types::BonsaiFileChange;
 use edenapi_types::CommitGraphEntry;
 use edenapi_types::CommitGraphRequest;
+use edenapi_types::CommitGraphSegmentParent;
+use edenapi_types::CommitGraphSegmentsEntry;
+use edenapi_types::CommitGraphSegmentsRequest;
 use edenapi_types::CommitHashLookupRequest;
 use edenapi_types::CommitHashLookupResponse;
 use edenapi_types::CommitHashToLocationResponse;
@@ -669,6 +672,60 @@ impl EdenApiHandler for GraphHandlerV2 {
     }
 }
 
+pub struct GraphSegmentsHandler;
+
+#[async_trait]
+impl EdenApiHandler for GraphSegmentsHandler {
+    type Request = CommitGraphSegmentsRequest;
+    type Response = CommitGraphSegmentsEntry;
+
+    const HTTP_METHOD: hyper::Method = hyper::Method::POST;
+    const API_METHOD: EdenApiMethod = EdenApiMethod::CommitGraphSegments;
+    const ENDPOINT: &'static str = "/commit/graph_segments";
+
+    async fn handler(
+        ectx: EdenApiContext<Self::PathExtractor, Self::QueryStringExtractor>,
+        request: Self::Request,
+    ) -> HandlerResult<'async_trait, Self::Response> {
+        let repo = ectx.repo();
+        let heads: Vec<_> = request
+            .heads
+            .into_iter()
+            .map(|hg_id| HgChangesetId::new(HgNodeHash::from(hg_id)))
+            .collect();
+        let common: Vec<_> = request
+            .common
+            .into_iter()
+            .map(|hg_id| HgChangesetId::new(HgNodeHash::from(hg_id)))
+            .collect();
+
+        let graph_segment_entries =
+            repo.graph_segments(common, heads)
+                .await?
+                .into_iter()
+                .map(|segment| {
+                    Ok(CommitGraphSegmentsEntry {
+                        head: HgId::from(segment.head.into_nodehash()),
+                        base: HgId::from(segment.base.into_nodehash()),
+                        length: segment.length,
+                        parents: segment
+                            .parents
+                            .into_iter()
+                            .map(|parent| CommitGraphSegmentParent {
+                                hgid: HgId::from(parent.hgid.into_nodehash()),
+                                location: parent.location.map(|location| {
+                                    location.map_descendant(|descendant| {
+                                        HgId::from(descendant.into_nodehash())
+                                    })
+                                }),
+                            })
+                            .collect(),
+                    })
+                });
+        Ok(stream::iter(graph_segment_entries).boxed())
+    }
+}
+
 pub struct CommitMutationsHandler;
 
 #[async_trait]

eden/mononoke/edenapi_service/src/handlers/mod.rs

@@ -111,6 +111,7 @@ pub enum EdenApiMethod {
     FetchSnapshot,
     AlterSnapshot,
     CommitGraphV2,
+    CommitGraphSegments,
     DownloadFile,
     CommitMutations,
     CommitTranslateId,
@@ -130,6 +131,7 @@ impl fmt::Display for EdenApiMethod {
             Self::CommitRevlogData => "commit_revlog_data",
             Self::CommitHashLookup => "commit_hash_lookup",
             Self::CommitGraphV2 => "commit_graph_v2",
+            Self::CommitGraphSegments => "commit_graph_segments",
             Self::Clone => "clone",
             Self::Bookmarks => "bookmarks",
             Self::SetBookmark => "set_bookmark",
@@ -438,6 +440,7 @@ pub fn build_router(ctx: ServerContext) -> Router {
         Handlers::setup::<commit::FetchSnapshotHandler>(route);
         Handlers::setup::<commit::AlterSnapshotHandler>(route);
         Handlers::setup::<commit::GraphHandlerV2>(route);
+        Handlers::setup::<commit::GraphSegmentsHandler>(route);
         Handlers::setup::<files::DownloadFileHandler>(route);
         Handlers::setup::<commit::CommitMutationsHandler>(route);
         Handlers::setup::<commit::CommitTranslateId>(route);

eden/mononoke/edenapi_service/src/middleware/ods.rs

@@ -52,6 +52,7 @@ define_stats! {
     alter_snapshot_duration_ms: histogram(100, 0, 5000, Average, Sum, Count; P 50; P 75; P 95; P 99),
     commit_graph_duration_ms: histogram(100, 0, 5000, Average, Sum, Count; P 50; P 75; P 95; P 99),
     commit_graph_v2_duration_ms: histogram(100, 0, 5000, Average, Sum, Count; P 50; P 75; P 95; P 99),
+    commit_graph_segments_duration_ms: histogram(100, 0, 5000, Average, Sum, Count; P 50; P 75; P 95; P 99),
     download_file_duration_ms: histogram(100, 0, 5000, Average, Sum, Count; P 50; P 75; P 95; P 99),
     commit_mutations_duration_ms: histogram(100, 0, 5000, Average, Sum, Count; P 50; P 75; P 95; P 99),
     commit_translate_id_duration_ms: histogram(100, 0, 5000, Average, Sum, Count; P 50; P 75; P 95; P 99),
@@ -112,6 +113,7 @@ fn log_stats(state: &mut State, status: StatusCode) -> Option<()> {
                 FetchSnapshot => STATS::fetch_snapshot_duration_ms.add_value(dur_ms),
                 AlterSnapshot => STATS::alter_snapshot_duration_ms.add_value(dur_ms),
                 CommitGraphV2 => STATS::commit_graph_v2_duration_ms.add_value(dur_ms),
+                CommitGraphSegments => STATS::commit_graph_segments_duration_ms.add_value(dur_ms),
                 DownloadFile => STATS::download_file_duration_ms.add_value(dur_ms),
                 CommitMutations => STATS::commit_mutations_duration_ms.add_value(dur_ms),
                 CommitTranslateId => STATS::commit_translate_id_duration_ms.add_value(dur_ms),

eden/mononoke/mononoke_api_hg/src/repo.rs

@@ -90,6 +90,18 @@ pub struct HgRepoContext {
     repo: RepoContext,
 }
 
+pub struct HgChangesetSegment {
+    pub head: HgChangesetId,
+    pub base: HgChangesetId,
+    pub length: u64,
+    pub parents: Vec<HgChangesetSegmentParent>,
+}
+
+pub struct HgChangesetSegmentParent {
+    pub hgid: HgChangesetId,
+    pub location: Option<Location<HgChangesetId>>,
+}
+
 impl HgRepoContext {
     pub(crate) fn new(repo: RepoContext) -> Self {
         Self { repo }
@@ -823,6 +835,94 @@ impl HgRepoContext {
         Ok(len == public_phases.len())
     }
 
+    pub async fn graph_segments(
+        &self,
+        common: Vec<HgChangesetId>,
+        heads: Vec<HgChangesetId>,
+    ) -> Result<Vec<HgChangesetSegment>, MononokeError> {
+        let bonsai_common = self.convert_changeset_ids(common).await?;
+        let bonsai_heads = self.convert_changeset_ids(heads).await?;
+
+        let segments = self
+            .repo()
+            .repo()
+            .commit_graph()
+            .ancestors_difference_segments(self.ctx(), bonsai_heads, bonsai_common)
+            .await?;
+
+        let bonsai_hg_mapping = stream::iter(segments.clone())
+            .flat_map(|segment| {
+                stream::iter([segment.head, segment.base])
+                    .chain(stream::iter(segment.parents).map(|parent| parent.cs_id))
+            })
+            .chunks(100)
+            .then(move |chunk| async move {
+                let mapping = self
+                    .blob_repo()
+                    .get_hg_bonsai_mapping(self.ctx().clone(), chunk.to_vec())
+                    .await
+                    .context("error fetching hg bonsai mapping")?;
+                Ok::<_, Error>(mapping)
+            })
+            .try_collect::<Vec<Vec<(HgChangesetId, ChangesetId)>>>()
+            .await?
+            .into_iter()
+            .flatten()
+            .map(|(hgid, csid)| (csid, hgid))
+            .collect::<HashMap<_, _>>();
+
+        segments
+            .into_iter()
+            .map(|segment| {
+                Ok(HgChangesetSegment {
+                    head: *bonsai_hg_mapping.get(&segment.head).ok_or_else(|| {
+                        MononokeError::InvalidRequest(format!(
+                            "failed to find hg equivalent for segment head {}",
+                            segment.head
+                        ))
+                    })?,
+                    base: *bonsai_hg_mapping.get(&segment.base).ok_or_else(|| {
+                        MononokeError::InvalidRequest(format!(
+                            "failed to find hg equivalent for segment base {}",
+                            segment.base
+                        ))
+                    })?,
+                    length: segment.length,
+                    parents: segment
+                        .parents
+                        .into_iter()
+                        .map(|parent| {
+                            Ok(HgChangesetSegmentParent {
+                                hgid: *bonsai_hg_mapping.get(&parent.cs_id).ok_or_else(|| {
+                                    MononokeError::InvalidRequest(format!(
+                                        "failed to find hg equivalent for segment parent {}",
+                                        parent
+                                    ))
+                                })?,
+                                location: parent
+                                    .location
+                                    .map(|location| {
+                                        Ok::<_, Error>(Location {
+                                            descendant: *bonsai_hg_mapping.get(&location.head).ok_or_else(
+                                                || {
+                                                    MononokeError::InvalidRequest(format!(
+                                                        "failed to find hg equivalent for location head {}",
+                                                        location.head
+                                                    ))
+                                                },
+                                            )?,
+                                            distance: location.distance,
+                                        })
+                                    })
+                                    .transpose()?,
+                            })
+                        })
+                        .collect::<Result<_, MononokeError>>()?,
+                })
+            })
+            .collect::<Result<_, MononokeError>>()
+    }
+
     /// Return a mapping of commits to their parents that are in the segment of
     /// of the commit graph bounded by common and heads.
     ///

eden/scm/lib/edenapi/types/src/commit.rs

@@ -128,6 +128,41 @@ pub struct CommitGraphRequest {
     pub heads: Vec<HgId>,
 }
 
+#[auto_wire]
+#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize)]
+#[cfg_attr(any(test, feature = "for-tests"), derive(Arbitrary))]
+pub struct CommitGraphSegmentsEntry {
+    #[id(1)]
+    pub head: HgId,
+    #[id(2)]
+    pub base: HgId,
+    #[id(3)]
+    pub length: u64,
+    #[id(4)]
+    pub parents: Vec<CommitGraphSegmentParent>,
+}
+
+#[auto_wire]
+#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize)]
+#[cfg_attr(any(test, feature = "for-tests"), derive(Arbitrary))]
+pub struct CommitGraphSegmentParent {
+    #[id(1)]
+    pub hgid: HgId,
+    #[id(2)]
+    pub location: Option<Location<HgId>>,
+}
+
+#[auto_wire]
+#[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
+#[derive(Serialize, Deserialize)]
+#[cfg_attr(any(test, feature = "for-tests"), derive(Arbitrary))]
+pub struct CommitGraphSegmentsRequest {
+    #[id(1)]
+    pub common: Vec<HgId>,
+    #[id(2)]
+    pub heads: Vec<HgId>,
+}
+
 /// The list of Mercurial commit identifiers for which we want the commit data to be returned.
 #[derive(Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
 #[derive(Serialize, Deserialize)]

eden/scm/lib/edenapi/types/src/lib.rs

@@ -75,6 +75,9 @@ pub use crate::commit::BonsaiChangesetContent;
 pub use crate::commit::BonsaiFileChange;
 pub use crate::commit::CommitGraphEntry;
 pub use crate::commit::CommitGraphRequest;
+pub use crate::commit::CommitGraphSegmentParent;
+pub use crate::commit::CommitGraphSegmentsEntry;
+pub use crate::commit::CommitGraphSegmentsRequest;
 pub use crate::commit::CommitHashLookupRequest;
 pub use crate::commit::CommitHashLookupResponse;
 pub use crate::commit::CommitHashToLocationRequestBatch;