feat(mev-rs): Refactor and Bench Validator Registry State Extension

mev-rs/Cargo.toml

@@ -4,8 +4,6 @@ version.workspace = true
 edition = "2021"
 license = "MIT OR Apache-2.0"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [features]
 default = ["serde", "builder-api", "relay-api"]
 builder-api = ["api"]
@@ -34,3 +32,8 @@ beacon-api-client = { workspace = true, optional = true }
 
 [dev-dependencies]
 rand = "0.8.5"
+criterion = { version = "0.5", features = ["html_reports"] }
+
+[[bench]]
+name = "registry"
+harness = false

mev-rs/benches/registry.rs

@@ -0,0 +1,41 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+
+use beacon_api_client::{ValidatorStatus, ValidatorSummary};
+use mev_rs::validator_registry::*;
+
+fn extend_summaries_grouped(v: Vec<ValidatorSummary>) {
+    let mut state = State::default();
+    for summary in v.into_iter() {
+        let public_key = summary.validator.public_key.clone();
+        state.pubkeys_by_index.insert(summary.index, public_key.clone());
+        state.validators.insert(public_key, summary);
+    }
+}
+
+fn extend_summaries_distributive(v: Vec<ValidatorSummary>) {
+    let mut state = State::default();
+    let _ = state.extend_summaries(v);
+}
+
+fn generate_dummy_summaries(n: usize) -> Vec<ValidatorSummary> {
+    (0..n)
+        .map(|x| ValidatorSummary {
+            index: x,
+            validator: Default::default(),
+            balance: 0,
+            status: ValidatorStatus::Active,
+        })
+        .collect()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    c.bench_function("grouped iterator", |b| {
+        b.iter(|| extend_summaries_grouped(black_box(generate_dummy_summaries(10))))
+    });
+    c.bench_function("distributive iterator", |b| {
+        b.iter(|| extend_summaries_distributive(black_box(generate_dummy_summaries(10))))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

mev-rs/src/lib.rs

@@ -9,7 +9,7 @@ pub mod relay;
 pub mod serde;
 pub mod signing;
 pub mod types;
-mod validator_registry;
+pub mod validator_registry;
 
 pub use blinded_block_provider::BlindedBlockProvider;
 pub use blinded_block_relayer::BlindedBlockRelayer;

mev-rs/src/validator_registry.rs

@@ -83,8 +83,25 @@ pub struct State {
     // data from registered validators
     validator_preferences: HashMap<BlsPublicKey, SignedValidatorRegistration>,
     // data from consensus
-    validators: HashMap<BlsPublicKey, ValidatorSummary>,
-    pubkeys_by_index: HashMap<ValidatorIndex, BlsPublicKey>,
+    pub validators: HashMap<BlsPublicKey, ValidatorSummary>,
+    pub pubkeys_by_index: HashMap<ValidatorIndex, BlsPublicKey>,
+}
+
+impl State {
+    /// Extends the [State] list of [`ValidatorSummary`] items.
+    pub fn extend_summaries(&mut self, summaries: Vec<ValidatorSummary>) -> Result<(), Error> {
+        let pubkeys_by_index = summaries
+            .iter()
+            .map(|summary| (summary.index, summary.validator.public_key.clone()))
+            .collect::<Vec<_>>();
+        let validators = summaries
+            .into_iter()
+            .map(|summary| (summary.validator.public_key.clone(), summary))
+            .collect::<Vec<_>>();
+        self.pubkeys_by_index.extend(pubkeys_by_index);
+        self.validators.extend(validators);
+        Ok(())
+    }
 }
 
 // Maintains validators we are aware of
@@ -100,17 +117,11 @@ impl ValidatorRegistry {
         Self { client, slots_per_epoch, state }
     }
 
-    // TODO: load more efficiently
     pub async fn on_epoch(&self, epoch: Epoch) -> Result<(), Error> {
         let slot = epoch * self.slots_per_epoch;
         let summaries = self.client.get_validators(StateId::Slot(slot), &[], &[]).await?;
         let mut state = self.state.write();
-        for summary in summaries.into_iter() {
-            let public_key = summary.validator.public_key.clone();
-            state.pubkeys_by_index.insert(summary.index, public_key.clone());
-            state.validators.insert(public_key, summary);
-        }
-        Ok(())
+        state.extend_summaries(summaries)
     }
 
     // Return the BLS public key for the validator's `index`, reflecting the index
@@ -166,7 +177,7 @@ impl ValidatorRegistry {
             let status =
                 determine_validator_registration_status(message.timestamp, latest_timestamp);
             if matches!(status, ValidatorRegistrationStatus::Outdated) {
-                return Err(Error::OutdatedRegistration(message.clone(), latest_timestamp))
+                return Err(Error::OutdatedRegistration(message.clone(), latest_timestamp));
             }
             status
         } else {