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 {
+    pub fn kzg(mut self, kzg: Option<Arc<Kzg>>) -> Self {
-        self.trusted_setup = Some(trusted_setup);
+        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)
+        let kzg = spec.deneb_fork_epoch.map(|_| KZG.clone());
-            .map_err(|e| format!("Unable to read trusted setup file: {}", e))
-            .unwrap();
 
         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)
+    let kzg_opt = spec.deneb_fork_epoch.map(|_| KZG.clone());
-        .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");
 
     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,
     )
 }
 

beacon_node/beacon_chain/tests/store_tests.rs

@@ -7,15 +7,13 @@ use beacon_chain::data_availability_checker::AvailableBlock;
 use beacon_chain::schema_change::migrate_schema;
 use beacon_chain::test_utils::{
     mock_execution_layer_from_parts, test_spec, AttestationStrategy, BeaconChainHarness,
-    BlockStrategy, DiskHarnessType,
+    BlockStrategy, DiskHarnessType, KZG,
 };
 use beacon_chain::{
     data_availability_checker::MaybeAvailableBlock, historical_blocks::HistoricalBlockError,
     migrate::MigratorConfig, BeaconChain, BeaconChainError, BeaconChainTypes, BeaconSnapshot,
     BlockError, ChainConfig, NotifyExecutionLayer, ServerSentEventHandler, WhenSlotSkipped,
 };
-use eth2_network_config::TRUSTED_SETUP_BYTES;
-use kzg::TrustedSetup;
 use lazy_static::lazy_static;
 use logging::test_logger;
 use maplit::hashset;
@@ -2418,9 +2416,7 @@ async fn weak_subjectivity_sync_test(slots: Vec<Slot>, checkpoint_slot: Slot) {
     let store = get_store(&temp2);
     let spec = test_spec::<E>();
     let seconds_per_slot = spec.seconds_per_slot;
-    let trusted_setup: TrustedSetup = serde_json::from_reader(TRUSTED_SETUP_BYTES)
+    let kzg = spec.deneb_fork_epoch.map(|_| KZG.clone());
-        .map_err(|e| println!("Unable to read trusted setup file: {}", e))
-        .unwrap();
 
     let mock =
         mock_execution_layer_from_parts(&harness.spec, harness.runtime.task_executor.clone());
@@ -2457,7 +2453,7 @@ async fn weak_subjectivity_sync_test(slots: Vec<Slot>, checkpoint_slot: Slot) {
             1,
         )))
         .execution_layer(Some(mock.el))
-        .trusted_setup(trusted_setup)
+        .kzg(kzg)
         .build()
         .expect("should build");
 

beacon_node/client/src/builder.rs

@@ -606,7 +606,12 @@ where
         };
 
         let beacon_chain_builder = if let Some(trusted_setup) = config.trusted_setup {
-            beacon_chain_builder.trusted_setup(trusted_setup)
+            let kzg = trusted_setup
+                .try_into()
+                .map(Arc::new)
+                .map(Some)
+                .map_err(|e| format!("Failed to load trusted setup: {:?}", e))?;
+            beacon_chain_builder.kzg(kzg)
         } else {
             beacon_chain_builder
         };

beacon_node/execution_layer/src/test_utils/execution_block_generator.rs

@@ -155,7 +155,7 @@ impl<E: EthSpec> ExecutionBlockGenerator<E> {
         shanghai_time: Option<u64>,
         cancun_time: Option<u64>,
         prague_time: Option<u64>,
-        kzg: Option<Kzg>,
+        kzg: Option<Arc<Kzg>>,
     ) -> Self {
         let mut gen = Self {
             head_block: <_>::default(),
@@ -172,7 +172,7 @@ impl<E: EthSpec> ExecutionBlockGenerator<E> {
             cancun_time,
             prague_time,
             blobs_bundles: <_>::default(),
-            kzg: kzg.map(Arc::new),
+            kzg,
             rng: make_rng(),
         };
 

beacon_node/execution_layer/src/test_utils/mock_execution_layer.rs

@@ -43,7 +43,7 @@ impl<E: EthSpec> MockExecutionLayer<E> {
         prague_time: Option<u64>,
         jwt_key: Option<JwtKey>,
         spec: ChainSpec,
-        kzg: Option<Kzg>,
+        kzg: Option<Arc<Kzg>>,
     ) -> Self {
         let handle = executor.handle().unwrap();
 

beacon_node/execution_layer/src/test_utils/mod.rs

@@ -109,7 +109,7 @@ impl<E: EthSpec> MockServer<E> {
     pub fn new_with_config(
         handle: &runtime::Handle,
         config: MockExecutionConfig,
-        kzg: Option<Kzg>,
+        kzg: Option<Arc<Kzg>>,
     ) -> Self {
         let MockExecutionConfig {
             jwt_key,
@@ -193,7 +193,7 @@ impl<E: EthSpec> MockServer<E> {
         shanghai_time: Option<u64>,
         cancun_time: Option<u64>,
         prague_time: Option<u64>,
-        kzg: Option<Kzg>,
+        kzg: Option<Arc<Kzg>>,
     ) -> Self {
         Self::new_with_config(
             handle,

crypto/kzg/src/lib.rs

@@ -142,3 +142,11 @@ impl Kzg {
         .map_err(Into::into)
     }
 }
+
+impl TryFrom<TrustedSetup> for Kzg {
+    type Error = Error;
+
+    fn try_from(trusted_setup: TrustedSetup) -> Result<Self, Self::Error> {
+        Kzg::new_from_trusted_setup(trusted_setup)
+    }
+}