Add cli flag to enable sampling and disable by default. (#6209)

sigp · Aug 1, 2024 · 37dd0ea · 37dd0ea
1 parent b148c4b
commit 37dd0ea
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 2 deletions.
diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -6960,8 +6960,10 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
     /// Returns true if we should issue a sampling request for this block
     /// TODO(das): check if the block is still within the da_window
     pub fn should_sample_slot(&self, slot: Slot) -> bool {
-        self.spec
-            .is_peer_das_enabled_for_epoch(slot.epoch(T::EthSpec::slots_per_epoch()))
+        self.config.enable_sampling
+            && self
+                .spec
+                .is_peer_das_enabled_for_epoch(slot.epoch(T::EthSpec::slots_per_epoch()))
     }
 
     pub fn logger(&self) -> &Logger {

diff --git a/beacon_node/beacon_chain/src/chain_config.rs b/beacon_node/beacon_chain/src/chain_config.rs
@@ -86,6 +86,8 @@ pub struct ChainConfig {
     pub enable_light_client_server: bool,
     /// Enable malicious PeerDAS mode where node withholds data columns when publishing a block
     pub malicious_withhold_count: usize,
+    /// Enable peer sampling on blocks.
+    pub enable_sampling: bool,
 }
 
 impl Default for ChainConfig {
@@ -118,6 +120,7 @@ impl Default for ChainConfig {
             epochs_per_migration: crate::migrate::DEFAULT_EPOCHS_PER_MIGRATION,
             enable_light_client_server: false,
             malicious_withhold_count: 0,
+            enable_sampling: false,
         }
     }
 }

diff --git a/beacon_node/src/cli.rs b/beacon_node/src/cli.rs
@@ -74,6 +74,15 @@ pub fn cli_app() -> Command {
                 .hide(true)
                 .display_order(0)
         )
+        .arg(
+            Arg::new("enable-sampling")
+                .long("enable-sampling")
+                .action(ArgAction::SetTrue)
+                .help_heading(FLAG_HEADER)
+                .help("Enable peer sampling on data columns. Disabled by default.")
+                .hide(true)
+                .display_order(0)
+        )
         .arg(
             Arg::new("subscribe-all-subnets")
                 .long("subscribe-all-subnets")

diff --git a/beacon_node/src/config.rs b/beacon_node/src/config.rs
@@ -181,6 +181,10 @@ pub fn get_config<E: EthSpec>(
         client_config.chain.shuffling_cache_size = cache_size;
     }
 
+    if cli_args.get_flag("enable-sampling") {
+        client_config.chain.enable_sampling = true;
+    }
+
     /*
      * Prometheus metrics HTTP server
      */

diff --git a/lighthouse/tests/beacon_node.rs b/lighthouse/tests/beacon_node.rs
@@ -832,6 +832,19 @@ fn network_subscribe_all_data_column_subnets_flag() {
         .with_config(|config| assert!(config.network.subscribe_all_data_column_subnets));
 }
 #[test]
+fn network_enable_sampling_flag() {
+    CommandLineTest::new()
+        .flag("enable-sampling", None)
+        .run_with_zero_port()
+        .with_config(|config| assert!(config.chain.enable_sampling));
+}
+#[test]
+fn network_enable_sampling_flag_default() {
+    CommandLineTest::new()
+        .run_with_zero_port()
+        .with_config(|config| assert!(!config.chain.enable_sampling));
+}
+#[test]
 fn network_subscribe_all_subnets_flag() {
     CommandLineTest::new()
         .flag("subscribe-all-subnets", None)