mononoke: add --all-types to backfill_derive_data single

Summary: It's useful to derive all enabled derived data at once

Reviewed By: krallin

Differential Revision: D22336338

fbshipit-source-id: 54bc27ab2c23c175913fc02e6bf05d18a54c249c

eden/mononoke/cmds/backfill_derived_data/main.rs

@@ -63,6 +63,7 @@ define_stats! {
     derivation_time_ms: dynamic_timeseries("{}.derivation_time_ms", (reponame: String); Average, Sum),
 }
 
+const ARG_ALL_TYPES: &str = "all-types";
 const ARG_DERIVED_DATA_TYPE: &str = "derived-data-type";
 const ARG_DRY_RUN: &str = "dry-run";
 const ARG_OUT_FILENAME: &str = "out-filename";
@@ -158,7 +159,7 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
                 .about("tail public commits and fill derived data")
                 .arg(
                     Arg::with_name(ARG_DERIVED_DATA_TYPE)
-                        .required(true)
+                        .required(false)
                         .multiple(true)
                         .index(1)
                         .possible_values(POSSIBLE_DERIVED_TYPES)
@@ -181,17 +182,25 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
             SubCommand::with_name(SUBCOMMAND_SINGLE)
                 .about("backfill single changeset (mainly for performance testing purposes)")
                 .arg(
-                    Arg::with_name(ARG_DERIVED_DATA_TYPE)
+                    Arg::with_name(ARG_ALL_TYPES)
-                        .required(true)
+                        .long(ARG_ALL_TYPES)
-                        .index(1)
+                        .required(false)
-                        .possible_values(POSSIBLE_DERIVED_TYPES)
+                        .takes_value(false)
-                        .help("derived data type for which backfill will be run"),
+                        .help("derive all derived data types enabled for this repo"),
                 )
                 .arg(
                     Arg::with_name(ARG_CHANGESET)
                         .required(true)
-                        .index(2)
+                        .index(1)
                         .help("changeset by {hd|bonsai} hash or bookmark"),
+                )
+                .arg(
+                    Arg::with_name(ARG_DERIVED_DATA_TYPE)
+                        .required(false)
+                        .index(2)
+                        .conflicts_with(ARG_ALL_TYPES)
+                        .possible_values(POSSIBLE_DERIVED_TYPES)
+                        .help("derived data type for which backfill will be run"),
                 ),
         );
     let matches = app.get_matches();
@@ -339,17 +348,47 @@ async fn run_subcmd<'a>(
                 .value_of_lossy(ARG_CHANGESET)
                 .ok_or_else(|| format_err!("missing required argument: {}", ARG_CHANGESET))?
                 .to_string();
-            let derived_data_type = sub_m
+            let all = sub_m.is_present(ARG_ALL_TYPES);
-                .value_of(ARG_DERIVED_DATA_TYPE)
+            let derived_data_type = sub_m.value_of(ARG_DERIVED_DATA_TYPE);
-                .ok_or_else(|| format_err!("missing required argument: {}", ARG_DERIVED_DATA_TYPE))?
+            let (repo, types): (_, Vec<String>) = match (all, derived_data_type) {
-                .to_string();
+                (true, None) => {
-            let repo = open_repo_maybe_unredacted(fb, &logger, &matches, &derived_data_type)
+                    let repo = args::open_repo_unredacted(fb, logger, matches)
-                .compat()
+                        .compat()
-                .await?;
+                        .await?;
+                    let types = repo
+                        .get_derived_data_config()
+                        .derived_data_types
+                        .clone()
+                        .into_iter()
+                        .collect();
+                    (repo, types)
+                }
+                (false, Some(derived_data_type)) => {
+                    let repo =
+                        open_repo_maybe_unredacted(fb, &logger, &matches, &derived_data_type)
+                            .compat()
+                            .await?;
+                    (repo, vec![derived_data_type.to_string()])
+                }
+                (true, Some(_)) => {
+                    return Err(format_err!(
+                        "{} and {} can't be specified",
+                        ARG_ALL_TYPES,
+                        ARG_DERIVED_DATA_TYPE
+                    ));
+                }
+                (false, None) => {
+                    return Err(format_err!(
+                        "{} or {} should be specified",
+                        ARG_ALL_TYPES,
+                        ARG_DERIVED_DATA_TYPE
+                    ));
+                }
+            };
             let csid = helpers::csid_resolve(ctx.clone(), repo.clone(), hash_or_bookmark)
                 .compat()
                 .await?;
-            subcommand_single(&ctx, &repo, csid, &derived_data_type).await
+            subcommand_single(&ctx, &repo, csid, types).await
         }
         (name, _) => Err(format_err!("unhandled subcommand: {}", name)),
     }
@@ -540,25 +579,36 @@ async fn subcommand_single(
     ctx: &CoreContext,
     repo: &BlobRepo,
     csid: ChangesetId,
-    derived_data_type: &str,
+    derived_data_types: Vec<String>,
 ) -> Result<(), Error> {
     let repo = repo.dangerous_override(|_| Arc::new(DummyLease {}) as Arc<dyn LeaseOps>);
-    let derived_utils = derived_data_utils(repo.clone(), derived_data_type)?;
+    let mut derived_utils = vec![];
-    derived_utils.regenerate(&vec![csid]);
+    for ty in derived_data_types {
-    derived_utils
+        let utils = derived_data_utils(repo.clone(), ty)?;
-        .derive(ctx.clone(), repo, csid)
+        utils.regenerate(&vec![csid]);
-        .timed({
+        derived_utils.push(utils);
-            cloned!(ctx);
+    }
-            move |stats, result| {
+    stream::iter(derived_utils)
-                info!(
+        .map(Ok)
-                    ctx.logger(),
+        .try_for_each_concurrent(100, |derived_utils| {
-                    "derived in {:?}: {:?}", stats.completion_time, result
+            derived_utils
-                );
+                .derive(ctx.clone(), repo.clone(), csid)
-                Ok(())
+                .timed({
-            }
+                    cloned!(ctx);
+                    move |stats, result| {
+                        info!(
+                            ctx.logger(),
+                            "derived {} in {:?}: {:?}",
+                            derived_utils.name(),
+                            stats.completion_time,
+                            result
+                        );
+                        Ok(())
+                    }
+                })
+                .map(|_| ())
+                .compat()
         })
-        .map(|_| ())
-        .compat()
         .await
 }
 
@@ -602,10 +652,22 @@ mod tests {
         let counting_blobstore = counting_blobstore.unwrap();
 
         let master = resolve_cs_id(&ctx, &repo, "master").await?;
-        subcommand_single(&ctx, &repo, master, RootUnodeManifestId::NAME).await?;
+        subcommand_single(
+            &ctx,
+            &repo,
+            master,
+            vec![RootUnodeManifestId::NAME.to_string()],
+        )
+        .await?;
 
         let writes_count = counting_blobstore.writes_count();
-        subcommand_single(&ctx, &repo, master, RootUnodeManifestId::NAME).await?;
+        subcommand_single(
+            &ctx,
+            &repo,
+            master,
+            vec![RootUnodeManifestId::NAME.to_string()],
+        )
+        .await?;
         assert!(counting_blobstore.writes_count() > writes_count);
         Ok(())
     }

eden/mononoke/tests/integration/test-backfill-derived-data.t

@@ -33,3 +33,10 @@ backfill derived data
 
   $ mononoke_admin --log-level ERROR derived-data exists "$DERIVED_DATA_TYPE" master_bookmark
   Derived: c3384961b16276f2db77df9d7c874bbe981cf0525bd6f84a502f919044f2dabd
+
+  $ backfill_derived_data single c3384961b16276f2db77df9d7c874bbe981cf0525bd6f84a502f919044f2dabd "$DERIVED_DATA_TYPE"
+  * using repo "repo" repoid RepositoryId(0) (glob)
+  * changeset resolved as: * (glob)
+  * derived fsnodes in * (glob)
+  $ backfill_derived_data single c3384961b16276f2db77df9d7c874bbe981cf0525bd6f84a502f919044f2dabd --all-types 2>&1 | grep derived | wc -l
+  8