Only load Kzg in tests if necessary and only load it once (#5555)

* Only load KZG once if necessary  in tests.

beacon_node/beacon_chain/src/builder.rs

@@ -25,7 +25,7 @@ use eth1::Config as Eth1Config;
 use execution_layer::ExecutionLayer;
 use fork_choice::{ForkChoice, ResetPayloadStatuses};
 use futures::channel::mpsc::Sender;
-use kzg::{Kzg, TrustedSetup};
+use kzg::Kzg;
 use operation_pool::{OperationPool, PersistedOperationPool};
 use parking_lot::{Mutex, RwLock};
 use proto_array::{DisallowedReOrgOffsets, ReOrgThreshold};
@@ -101,7 +101,7 @@ pub struct BeaconChainBuilder<T: BeaconChainTypes> {
     // Pending I/O batch that is constructed during building and should be executed atomically
     // alongside `PersistedBeaconChain` storage when `BeaconChainBuilder::build` is called.
     pending_io_batch: Vec<KeyValueStoreOp>,
-    trusted_setup: Option<TrustedSetup>,
+    kzg: Option<Arc<Kzg>>,
     task_executor: Option<TaskExecutor>,
     validator_monitor_config: Option<ValidatorMonitorConfig>,
 }
@@ -142,7 +142,7 @@ where
             graffiti: Graffiti::default(),
             slasher: None,
             pending_io_batch: vec![],
-            trusted_setup: None,
+            kzg: None,
             task_executor: None,
             validator_monitor_config: None,
         }
@@ -669,8 +669,8 @@ where
         self
     }
 
-    pub fn trusted_setup(mut self, trusted_setup: TrustedSetup) -> Self {
-        self.trusted_setup = Some(trusted_setup);
+    pub fn kzg(mut self, kzg: Option<Arc<Kzg>>) -> Self {
+        self.kzg = kzg;
         self
     }
 
@@ -718,15 +718,6 @@ where
             slot_clock.now().ok_or("Unable to read slot")?
         };
 
-        let kzg = if let Some(trusted_setup) = self.trusted_setup {
-            let kzg = Kzg::new_from_trusted_setup(trusted_setup)
-                .map_err(|e| format!("Failed to load trusted setup: {:?}", e))?;
-            let kzg_arc = Arc::new(kzg);
-            Some(kzg_arc)
-        } else {
-            None
-        };
-
         let initial_head_block_root = fork_choice
             .get_head(current_slot, &self.spec)
             .map_err(|e| format!("Unable to get fork choice head: {:?}", e))?;
@@ -967,10 +958,10 @@ where
             validator_monitor: RwLock::new(validator_monitor),
             genesis_backfill_slot,
             data_availability_checker: Arc::new(
-                DataAvailabilityChecker::new(slot_clock, kzg.clone(), store, &log, self.spec)
+                DataAvailabilityChecker::new(slot_clock, self.kzg.clone(), store, &log, self.spec)
                     .map_err(|e| format!("Error initializing DataAvailabiltyChecker: {:?}", e))?,
             ),
-            kzg,
+            kzg: self.kzg.clone(),
             block_production_state: Arc::new(Mutex::new(None)),
         };
 

beacon_node/beacon_chain/src/lib.rs

@@ -86,7 +86,7 @@ pub use events::ServerSentEventHandler;
 pub use execution_layer::EngineState;
 pub use execution_payload::NotifyExecutionLayer;
 pub use fork_choice::{ExecutionStatus, ForkchoiceUpdateParameters};
-pub use kzg::TrustedSetup;
+pub use kzg::{Kzg, TrustedSetup};
 pub use metrics::scrape_for_metrics;
 pub use migrate::MigratorConfig;
 pub use parking_lot;

beacon_node/beacon_chain/src/test_utils.rs

@@ -31,6 +31,7 @@ use futures::channel::mpsc::Receiver;
 pub use genesis::{interop_genesis_state_with_eth1, DEFAULT_ETH1_BLOCK_HASH};
 use int_to_bytes::int_to_bytes32;
 use kzg::{Kzg, TrustedSetup};
+use lazy_static::lazy_static;
 use merkle_proof::MerkleTree;
 use operation_pool::ReceivedPreCapella;
 use parking_lot::Mutex;
@@ -76,6 +77,16 @@ pub const FORK_NAME_ENV_VAR: &str = "FORK_NAME";
 // a different value.
 pub const DEFAULT_TARGET_AGGREGATORS: u64 = u64::MAX;
 
+lazy_static! {
+    pub static ref KZG: Arc<Kzg> = {
+        let trusted_setup: TrustedSetup = serde_json::from_reader(TRUSTED_SETUP_BYTES)
+            .map_err(|e| format!("Unable to read trusted setup file: {}", e))
+            .expect("should have trusted setup");
+        let kzg = Kzg::new_from_trusted_setup(trusted_setup).expect("should create kzg");
+        Arc::new(kzg)
+    };
+}
+
 pub type BaseHarnessType<E, THotStore, TColdStore> =
     Witness<TestingSlotClock, CachingEth1Backend<E>, E, THotStore, TColdStore>;
 
@@ -505,16 +516,14 @@ where
         let validator_keypairs = self
             .validator_keypairs
             .expect("cannot build without validator keypairs");
-        let trusted_setup: TrustedSetup = serde_json::from_reader(TRUSTED_SETUP_BYTES)
-            .map_err(|e| format!("Unable to read trusted setup file: {}", e))
-            .unwrap();
+        let kzg = spec.deneb_fork_epoch.map(|_| KZG.clone());
 
         let validator_monitor_config = self.validator_monitor_config.unwrap_or_default();
 
         let chain_config = self.chain_config.unwrap_or_default();
         let mut builder = BeaconChainBuilder::new(self.eth_spec_instance)
             .logger(log.clone())
-            .custom_spec(spec)
+            .custom_spec(spec.clone())
             .store(self.store.expect("cannot build without store"))
             .store_migrator_config(
                 MigratorConfig::default()
@@ -532,7 +541,7 @@ where
                 5,
             )))
             .validator_monitor_config(validator_monitor_config)
-            .trusted_setup(trusted_setup);
+            .kzg(kzg);
 
         builder = if let Some(mutator) = self.initial_mutator {
             mutator(builder)
@@ -587,10 +596,7 @@ pub fn mock_execution_layer_from_parts<E: EthSpec>(
         HARNESS_GENESIS_TIME + spec.seconds_per_slot * E::slots_per_epoch() * epoch.as_u64()
     });
 
-    let trusted_setup: TrustedSetup = serde_json::from_reader(TRUSTED_SETUP_BYTES)
-        .map_err(|e| format!("Unable to read trusted setup file: {}", e))
-        .expect("should have trusted setup");
-    let kzg = Kzg::new_from_trusted_setup(trusted_setup).expect("should create kzg");
+    let kzg_opt = spec.deneb_fork_epoch.map(|_| KZG.clone());
 
     MockExecutionLayer::new(
         task_executor,
@@ -600,7 +606,7 @@ pub fn mock_execution_layer_from_parts<E: EthSpec>(
         prague_time,
         Some(JwtKey::from_slice(&DEFAULT_JWT_SECRET).unwrap()),
         spec.clone(),
-        Some(kzg),
+        kzg_opt,
     )
 }