From 2bda7e3d1487719db7c5b5e6e09ae53ccb03aa2e Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sun, 27 Jan 2019 15:45:29 +1100
Subject: [PATCH] Add `BeaconChain` benchmarking

---
 .../beacon_chain/test_harness/Cargo.toml      |  7 ++
 .../test_harness/benches/state_transition.rs  | 64 +++++++++++++++++++
 .../test_harness/src/beacon_chain_harness.rs  | 16 +++--
 3 files changed, 80 insertions(+), 7 deletions(-)
 create mode 100644 beacon_node/beacon_chain/test_harness/benches/state_transition.rs

diff --git a/beacon_node/beacon_chain/test_harness/Cargo.toml b/beacon_node/beacon_chain/test_harness/Cargo.toml
index 6972e4df2d..72fdf7164d 100644
--- a/beacon_node/beacon_chain/test_harness/Cargo.toml
+++ b/beacon_node/beacon_chain/test_harness/Cargo.toml
@@ -4,6 +4,13 @@ version = "0.1.0"
 authors = ["Paul Hauner <paul@paulhauner.com>"]
 edition = "2018"
 
+[[bench]]
+name = "state_transition"
+harness = false
+
+[dev-dependencies]
+criterion = "0.2"
+
 [dependencies]
 beacon_chain = { path = "../../beacon_chain" }
 block_producer = { path = "../../../eth2/block_producer" }
diff --git a/beacon_node/beacon_chain/test_harness/benches/state_transition.rs b/beacon_node/beacon_chain/test_harness/benches/state_transition.rs
new file mode 100644
index 0000000000..d78a177208
--- /dev/null
+++ b/beacon_node/beacon_chain/test_harness/benches/state_transition.rs
@@ -0,0 +1,64 @@
+use criterion::Criterion;
+use criterion::{criterion_group, criterion_main};
+use test_harness::BeaconChainHarness;
+use types::ChainSpec;
+
+fn mid_epoch_state_transition(c: &mut Criterion) {
+    let validator_count = 2;
+    let mut rig = BeaconChainHarness::new(ChainSpec::foundation(), validator_count);
+
+    let two_and_half_epochs = (rig.spec.epoch_length * 2) + (rig.spec.epoch_length / 2);
+
+    for _ in 0..two_and_half_epochs {
+        rig.advance_chain_with_block();
+    }
+
+    let block = rig.advance_chain_without_block();
+    let state = rig.beacon_chain.canonical_head().beacon_state.clone();
+
+    c.bench_function("mid-epoch state transition 10k validators", move |b| {
+        let block = block.clone();
+        let state = state.clone();
+        b.iter(|| {
+            rig.beacon_chain
+                .state_transition(state.clone(), &block.clone())
+        })
+    });
+}
+
+fn epoch_boundary_state_transition(c: &mut Criterion) {
+    let validator_count = 10_000;
+    let mut rig = BeaconChainHarness::new(ChainSpec::foundation(), validator_count);
+
+    let three_epochs = rig.spec.epoch_length * 3;
+
+    for _ in 0..(three_epochs - 1) {
+        rig.advance_chain_with_block();
+    }
+
+    let state = rig.beacon_chain.canonical_head().beacon_state.clone();
+    assert_eq!(
+        state.slot % rig.spec.epoch_length,
+        rig.spec.epoch_length - 1,
+    );
+    let block = rig.advance_chain_without_block();
+
+    c.bench_function("epoch boundary state transition 10k validators", move |b| {
+        let block = block.clone();
+        let state = state.clone();
+        b.iter(|| {
+            let state = rig
+                .beacon_chain
+                .state_transition(state.clone(), &block.clone())
+                .unwrap();
+            assert_eq!(state.slot % rig.spec.epoch_length, 0);
+        })
+    });
+}
+
+criterion_group!(
+    benches,
+    mid_epoch_state_transition,
+    epoch_boundary_state_transition
+);
+criterion_main!(benches);
diff --git a/beacon_node/beacon_chain/test_harness/src/beacon_chain_harness.rs b/beacon_node/beacon_chain/test_harness/src/beacon_chain_harness.rs
index 01318727b1..2f4c11aef0 100644
--- a/beacon_node/beacon_chain/test_harness/src/beacon_chain_harness.rs
+++ b/beacon_node/beacon_chain/test_harness/src/beacon_chain_harness.rs
@@ -1,12 +1,10 @@
 use super::TestValidator;
 pub use beacon_chain::dump::{Error as DumpError, SlotDump};
 use beacon_chain::BeaconChain;
-use block_producer::BeaconNode;
 use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
     MemoryDB,
 };
-use serde_json::Result as SerdeResult;
 use slot_clock::TestingSlotClock;
 use std::fs::File;
 use std::io::prelude::*;
@@ -14,11 +12,11 @@ use std::sync::Arc;
 use types::{BeaconBlock, ChainSpec, Keypair, Validator};
 
 pub struct BeaconChainHarness {
-    db: Arc<MemoryDB>,
-    beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock>>,
-    block_store: Arc<BeaconBlockStore<MemoryDB>>,
-    state_store: Arc<BeaconStateStore<MemoryDB>>,
-    validators: Vec<TestValidator>,
+    pub db: Arc<MemoryDB>,
+    pub beacon_chain: Arc<BeaconChain<MemoryDB, TestingSlotClock>>,
+    pub block_store: Arc<BeaconBlockStore<MemoryDB>>,
+    pub state_store: Arc<BeaconStateStore<MemoryDB>>,
+    pub validators: Vec<TestValidator>,
     pub spec: ChainSpec,
 }
 
@@ -76,6 +74,10 @@ impl BeaconChainHarness {
         }
     }
 
+    pub fn advance_chain_without_block(&mut self) -> BeaconBlock {
+        self.produce_next_slot()
+    }
+
     pub fn advance_chain_with_block(&mut self) {
         let block = self.produce_next_slot();
         self.beacon_chain.process_block(block).unwrap();