Allow required alias codemod to be applied to a rollout range

Reviewed By: gordyf

Differential Revision: D69144209

fbshipit-source-id: cf5cd4d933fadfbe84159266153459e3e0827de2

Author: captbaritone

compiler/crates/common/src/feature_flags.rs

@@ -16,6 +16,7 @@ use schemars::JsonSchema;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::rollout::RolloutRange;
 use crate::Rollout;
 
 #[derive(Default, Debug, Serialize, Deserialize, Clone, JsonSchema)]
@@ -183,6 +184,9 @@ pub enum FeatureFlag {
 
     /// Partially enabled: used for gradual rollout of the feature
     Rollout { rollout: Rollout },
+
+    /// Partially enabled: used for gradual rollout of the feature
+    RolloutRange { rollout: RolloutRange },
 }
 
 impl FeatureFlag {
@@ -191,6 +195,7 @@ impl FeatureFlag {
             FeatureFlag::Enabled => true,
             FeatureFlag::Limited { allowlist } => allowlist.contains(&name),
             FeatureFlag::Rollout { rollout } => rollout.check(name.lookup()),
+            FeatureFlag::RolloutRange { rollout } => rollout.check(name.lookup()),
             FeatureFlag::Disabled => false,
         }
     }
@@ -211,6 +216,7 @@ impl Display for FeatureFlag {
                 f.write_str(&items.join(", "))
             }
             FeatureFlag::Rollout { rollout } => write!(f, "Rollout: {:#?}", rollout),
+            FeatureFlag::RolloutRange { rollout } => write!(f, "RolloutRange: {:#?}", rollout),
         }
     }
 }

compiler/crates/common/src/lib.rs

@@ -63,5 +63,6 @@ pub use perf_logger::PerfLogEvent;
 pub use perf_logger::PerfLogger;
 pub use pointer_address::PointerAddress;
 pub use rollout::Rollout;
+pub use rollout::RolloutRange;
 pub use span::Span;
 pub use text_source::TextSource;

compiler/crates/common/src/rollout.rs

@@ -31,3 +31,23 @@ impl Rollout {
         }
     }
 }
+
+/// A utility to enable gradual rollout of large codegen changes. Allows you to
+/// specify a range of percentages to rollout.
+#[derive(Debug, Serialize, Deserialize, Clone, Copy, JsonSchema)]
+pub struct RolloutRange {
+    pub start: u8,
+    pub end: u8,
+}
+
+impl RolloutRange {
+    /// Checks some key deterministically and passes on average the given
+    /// percentage of the rollout.
+    /// A typical key to pass in could be the fragment or operation name.
+    pub fn check(&self, key: impl AsRef<[u8]>) -> bool {
+        let hash = Md5::digest(key.as_ref());
+        let hash: u16 = ((hash[1] as u16) << 8) | (hash[0] as u16);
+        let percent = hash % 100;
+        (percent) <= (self.end as u16) && (percent) >= (self.start as u16)
+    }
+}

compiler/crates/relay-codegen/src/build_ast.rs

@@ -12,7 +12,6 @@ use ::intern::string_key::Intern;
 use ::intern::string_key::StringKey;
 use ::intern::Lookup;
 use common::DirectiveName;
-use common::FeatureFlag;
 use common::NamedItem;
 use common::ObjectName;
 use common::WithLocation;
@@ -2394,17 +2393,11 @@ impl<'schema, 'builder, 'config> CodegenBuilder<'schema, 'builder, 'config> {
             kind: Primitive::String(CODEGEN_CONSTANTS.module_import),
         };
 
-        let should_use_reader_module_imports =
-            match &self.project_config.feature_flags.use_reader_module_imports {
-                FeatureFlag::Enabled => true,
-                FeatureFlag::Disabled => false,
-                FeatureFlag::Limited {
-                    allowlist: fragment_names,
-                } => fragment_names.contains(&module_metadata.key),
-                FeatureFlag::Rollout { rollout } => {
-                    rollout.check(module_metadata.key.lookup().as_bytes())
-                }
-            };
+        let should_use_reader_module_imports = self
+            .project_config
+            .feature_flags
+            .use_reader_module_imports
+            .is_enabled_for(module_metadata.key);
 
         match self.variant {
             CodegenVariant::Reader => {

compiler/crates/relay-codemod/src/codemod.rs

@@ -12,6 +12,7 @@ use clap::Args;
 use clap::Subcommand;
 use common::FeatureFlag;
 use common::Rollout;
+use common::RolloutRange;
 use log::info;
 use lsp_types::CodeActionOrCommand;
 use lsp_types::TextEdit;
@@ -32,9 +33,11 @@ pub enum AvailableCodemod {
 
 #[derive(Args, Debug, Clone)]
 pub struct MarkDangerousConditionalFragmentSpreadsArgs {
-    /// If using a feature flag, specify the rollout percentage. If omitted, the codemod is fully enabled.
-    #[clap(long, short, value_parser=valid_percent)]
-    pub rollout_percentage: Option<u8>,
+    /// Specify a percentage of fragments to codemod. If a number is provided,
+    /// the first n percentage of fragments will be codemodded. If a range (`20-30`) is
+    /// provided, then fragments between the start and end of the range will be codemodded.
+    #[clap(long, short, value_parser=valid_percent, default_value = "100")]
+    pub rollout_percentage: FeatureFlag,
 }
 
 pub async fn run_codemod(
@@ -46,12 +49,7 @@ pub async fn run_codemod(
         .iter()
         .flat_map(|programs| match &codemod {
             AvailableCodemod::MarkDangerousConditionalFragmentSpreads(opts) => {
-                match fragment_alias_directive(
-                    &programs.source,
-                    &FeatureFlag::Rollout {
-                        rollout: Rollout(opts.rollout_percentage),
-                    },
-                ) {
+                match fragment_alias_directive(&programs.source, &opts.rollout_percentage) {
                     Ok(_) => vec![],
                     Err(e) => e,
                 }
@@ -143,13 +141,33 @@ fn sort_changes(url: &Url, changes: &mut Vec<TextEdit>) -> Result<(), std::io::E
     Ok(())
 }
 
-fn valid_percent(s: &str) -> Result<u8, String> {
-    // turn s into a u8
-    let s = s.parse::<u8>().map_err(|_| "not a number".to_string())?;
-    // check if s is less than 100
-    if (0..=100).contains(&s) {
-        Ok(s)
+fn valid_percent(s: &str) -> Result<FeatureFlag, String> {
+    // If the string is a range of the form "x-y", where x and y are numbers, return the range
+    let parts: Vec<&str> = s.split('-').collect();
+    if parts.len() == 2 {
+        let start = parts[0].parse::<u8>().map_err(|_| {
+            "Expected the value on the left of the rollout range to be a number".to_string()
+        })?;
+        let end = parts[1].parse::<u8>().map_err(|_| {
+            "Expected the value on the right of the rollout range to be a number".to_string()
+        })?;
+        if (0..=100).contains(&start) && (0..=100).contains(&end) && start <= end {
+            Ok(FeatureFlag::RolloutRange {
+                rollout: RolloutRange { start, end },
+            })
+        } else {
+            Err("numbers must be between 0 and 100, inclusive, and the first number must be less than or equal to the second".to_string())
+        }
     } else {
-        Err("number must be between 0 and 100, inclusive".to_string())
+        // turn s into a u8
+        let s = s.parse::<u8>().map_err(|_| "not a number".to_string())?;
+        // check if s is less than 100
+        if (0..=100).contains(&s) {
+            Ok(FeatureFlag::Rollout {
+                rollout: Rollout(Some(s)),
+            })
+        } else {
+            Err("number must be between 0 and 100, inclusive".to_string())
+        }
     }
 }