From ae521ed84371280d99a4bdaf7801809bd40ae85f Mon Sep 17 00:00:00 2001
From: Rajiv Sharma <rajshar@meta.com>
Date: Wed, 12 Feb 2025 10:40:12 -0800
Subject: [PATCH] Allow Mononoke App to warm specific set of derived data type

Summary:
Right now, any application or service created using `MononokeApp` is capable of warming either:
- Git derived data types
- Hg derived data types
- All derived data types

Since there is no option that allows specifing a custom set of derived data types to be warmed, apps and services end up using the third option when in reality they just need a handful of derived data types. This diff incorporates a new variant, `BookmarkCacheDerivedData::SpecificTypes` which will allow for warming specific types.

Reviewed By: YousefSalama

Differential Revision: D69526247

fbshipit-source-id: 269bf1deb9a9594473c0ffe68fbc351ac1d38bdf
---
 .../bookmarks/warm_bookmarks_cache/lib.rs         | 11 +++++++++++
 eden/mononoke/cmdlib/environment/Cargo.toml       |  1 +
 eden/mononoke/cmdlib/environment/lib.rs           | 15 +++++----------
 eden/mononoke/cmdlib/mononoke_app/src/args/wbc.rs |  2 +-
 eden/mononoke/repo_factory/src/lib.rs             |  7 +++++++
 5 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/eden/mononoke/bookmarks/warm_bookmarks_cache/lib.rs b/eden/mononoke/bookmarks/warm_bookmarks_cache/lib.rs
index ce26965ef509e..eade75b6852db 100644
--- a/eden/mononoke/bookmarks/warm_bookmarks_cache/lib.rs
+++ b/eden/mononoke/bookmarks/warm_bookmarks_cache/lib.rs
@@ -212,6 +212,17 @@ impl WarmBookmarksCacheBuilder {
         Ok(())
     }
 
+    pub fn add_specific_types_warmers(
+        &mut self,
+        repo_derived_data: &ArcRepoDerivedData,
+        types: &[DerivableType],
+        phases: &ArcPhases,
+    ) -> Result<(), Error> {
+        self.add_derived_data_warmers(types, repo_derived_data)?;
+        self.add_public_phase_warmer(phases);
+        Ok(())
+    }
+
     fn add_derived_data_warmers<'a>(
         &mut self,
         types: impl IntoIterator<Item = &'a DerivableType>,
diff --git a/eden/mononoke/cmdlib/environment/Cargo.toml b/eden/mononoke/cmdlib/environment/Cargo.toml
index 0329489672322..9bb0c09b3dace 100644
--- a/eden/mononoke/cmdlib/environment/Cargo.toml
+++ b/eden/mononoke/cmdlib/environment/Cargo.toml
@@ -18,6 +18,7 @@ commit_graph_types = { version = "0.1.0", path = "../../repo_attributes/commit_g
 derived_data_remote = { version = "0.1.0", path = "../../derived_data/remote" }
 fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 megarepo_config = { version = "0.1.0", path = "../../megarepo_api/megarepo_config" }
+mononoke_types = { version = "0.1.0", path = "../../mononoke_types" }
 observability = { version = "0.1.0", path = "../../observability" }
 permission_checker = { version = "0.1.0", path = "../../permission_checker" }
 rendezvous = { version = "0.1.0", path = "../../common/rendezvous" }
diff --git a/eden/mononoke/cmdlib/environment/lib.rs b/eden/mononoke/cmdlib/environment/lib.rs
index a259c75bc7b9a..0940301f94de3 100644
--- a/eden/mononoke/cmdlib/environment/lib.rs
+++ b/eden/mononoke/cmdlib/environment/lib.rs
@@ -18,6 +18,7 @@ use commit_graph_types::environment::CommitGraphOptions;
 use derived_data_remote::RemoteDerivationOptions;
 use fbinit::FacebookInit;
 use megarepo_config::MononokeMegarepoConfigsOptions;
+use mononoke_types::DerivableType;
 use observability::ObservabilityContext;
 use permission_checker::AclProvider;
 use rendezvous::RendezVousOptions;
@@ -45,16 +46,7 @@ pub enum Caching {
     Disabled,
 }
 
-#[derive(
-    Copy,
-    Clone,
-    Debug,
-    ValueEnum,
-    EnumString,
-    strum::Display,
-    PartialEq,
-    Eq
-)]
+#[derive(Clone, Debug, ValueEnum, EnumString, strum::Display, PartialEq, Eq)]
 
 /// Which derived data types should the cache wait for before
 /// exposing the bookmark move to the users.
@@ -65,6 +57,9 @@ pub enum BookmarkCacheDerivedData {
     GitOnly,
     /// Wait for all derived data types - mainly used by Mononoke SCS Server.
     AllKinds,
+    /// Wait for specific derived data types.
+    #[value(skip)]
+    SpecificTypes(Vec<DerivableType>),
     /// Don't wait for any derived data - advance bookmarks as they move.
     NoDerivation,
 }
diff --git a/eden/mononoke/cmdlib/mononoke_app/src/args/wbc.rs b/eden/mononoke/cmdlib/mononoke_app/src/args/wbc.rs
index 4b2c2ef50e755..18d2c42596ca7 100644
--- a/eden/mononoke/cmdlib/mononoke_app/src/args/wbc.rs
+++ b/eden/mononoke/cmdlib/mononoke_app/src/args/wbc.rs
@@ -54,7 +54,7 @@ impl AppExtension for WarmBookmarksCacheExtension {
                 .as_ref()
                 .map(|smc_tier| BookmarkCacheAddress::SmcTier(smc_tier.to_string()))
         };
-        if let Some(derived_data) = args.enable_wbc_with {
+        if let Some(derived_data) = args.enable_wbc_with.clone() {
             // Enable_wbc_with enables local cache by default
             if env.bookmark_cache_options.cache_kind == BookmarkCacheKind::Disabled {
                 env.bookmark_cache_options.cache_kind = BookmarkCacheKind::Local;
diff --git a/eden/mononoke/repo_factory/src/lib.rs b/eden/mononoke/repo_factory/src/lib.rs
index c3f08634f5964..fd0c555ae744b 100644
--- a/eden/mononoke/repo_factory/src/lib.rs
+++ b/eden/mononoke/repo_factory/src/lib.rs
@@ -1675,6 +1675,13 @@ impl RepoFactory {
                     BookmarkCacheDerivedData::AllKinds => {
                         wbc_builder.add_all_warmers(repo_derived_data, phases)?;
                     }
+                    BookmarkCacheDerivedData::SpecificTypes(ref types) => {
+                        wbc_builder.add_specific_types_warmers(
+                            repo_derived_data,
+                            types.as_slice(),
+                            phases,
+                        )?;
+                    }
                     BookmarkCacheDerivedData::NoDerivation => {}
                 }