log completion_time_us for find_draft_ancestors

Summary: We log a sample with "msg": "Started finding draft ancestors", and we log a sample with "msg", "Found draft ancestors", but we never measure the time for it in the `completion_time_us` column. E.g.: https://fburl.com/scuba/mononoke_git_server/c2hnwrjr This diff fixes that Reviewed By: markbt Differential Revision: D66970659 fbshipit-source-id: 16202c3af978abb68d7294b0622266a452e09d88
facebook · Dec 9, 2024 · 0d3b9a6 · 0d3b9a6
1 parent 7dd10d6
commit 0d3b9a6
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 27 deletions.
diff --git a/eden/mononoke/features/repo_update_logger/Cargo.toml b/eden/mononoke/features/repo_update_logger/Cargo.toml
@@ -20,6 +20,7 @@ commit_graph = { version = "0.1.0", path = "../../repo_attributes/commit_graph/c
 context = { version = "0.1.0", path = "../../server/context" }
 ephemeral_blobstore = { version = "0.1.0", path = "../../blobstore/ephemeral_blobstore" }
 futures = { version = "0.3.30", features = ["async-await", "compat"] }
+futures_stats = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 git_source_of_truth = { version = "0.1.0", path = "../../git_source_of_truth" }
 gix-hash = "0.15.1"
 hostname = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }

diff --git a/eden/mononoke/features/repo_update_logger/src/commit_logger.rs b/eden/mononoke/features/repo_update_logger/src/commit_logger.rs
@@ -25,6 +25,7 @@ use ephemeral_blobstore::BubbleId;
 use futures::stream;
 use futures::StreamExt;
 use futures::TryStreamExt;
+use futures_stats::TimedTryFutureExt;
 use logger_ext::Loggable;
 use metaconfig_types::RepoConfigRef;
 #[cfg(fbcode_build)]
@@ -359,40 +360,45 @@ pub async fn find_draft_ancestors(
         .clone()
         .log_with_msg("Started finding draft ancestors", None);
 
-    let phases = repo.phases();
-    let mut queue = VecDeque::new();
-    let mut visited = HashSet::new();
-    let mut drafts = vec![];
-    queue.push_back(to_cs_id);
-    visited.insert(to_cs_id);
-
-    while let Some(cs_id) = queue.pop_front() {
-        let public = phases
-            .get_public(ctx, vec![cs_id], false /*ephemeral_derive*/)
-            .await?;
-
-        if public.contains(&cs_id) {
-            continue;
-        }
-        drafts.push(cs_id);
+    let (stats, drafts) = async move {
+        let phases = repo.phases();
+        let mut queue = VecDeque::new();
+        let mut visited = HashSet::new();
+        let mut drafts = vec![];
+        queue.push_back(to_cs_id);
+        visited.insert(to_cs_id);
+
+        while let Some(cs_id) = queue.pop_front() {
+            let public = phases
+                .get_public(ctx, vec![cs_id], false /*ephemeral_derive*/)
+                .await?;
+
+            if public.contains(&cs_id) {
+                continue;
+            }
+            drafts.push(cs_id);
 
-        let parents = repo.commit_graph().changeset_parents(ctx, cs_id).await?;
-        for p in parents {
-            if visited.insert(p) {
-                queue.push_back(p);
+            let parents = repo.commit_graph().changeset_parents(ctx, cs_id).await?;
+            for p in parents {
+                if visited.insert(p) {
+                    queue.push_back(p);
+                }
             }
         }
-    }
 
-    let drafts = stream::iter(drafts)
-        .map(Ok)
-        .map_ok(|cs_id| async move { cs_id.load(ctx, repo.repo_blobstore()).await })
-        .try_buffer_unordered(100)
-        .try_collect::<Vec<_>>()
-        .await?;
+        stream::iter(drafts)
+            .map(Ok)
+            .map_ok(|cs_id| async move { cs_id.load(ctx, repo.repo_blobstore()).await })
+            .try_buffer_unordered(100)
+            .try_collect::<Vec<_>>()
+            .await
+    }
+    .try_timed()
+    .await?;
 
     ctx.scuba()
         .clone()
+        .add_future_stats(&stats)
         .log_with_msg("Found draft ancestors", Some(format!("{}", drafts.len())));
     Ok(drafts)
 }