Add BenchingBeaconNode to chain tests

beacon_node/beacon_chain/tests/chain.rs

@@ -12,7 +12,7 @@ fn it_can_produce_blocks() {
     let blocks = rig.spec.epoch_length + 1;
 
     for _ in 0..blocks {
-        rig.produce_next_slot();
+        rig.advance_chain_with_block();
     }
     let dump = rig.chain_dump().expect("Chain dump failed.");
 

beacon_node/beacon_chain/tests/utils/benching_beacon_node.rs

@@ -0,0 +1,90 @@
+use beacon_chain::block_processing::{Error as ProcessingError, Outcome as ProcessingOutcome};
+use beacon_chain::{block_production::Error as BlockProductionError, BeaconChain};
+use block_producer::{
+    BeaconNode as BeaconBlockNode, BeaconNodeError as BeaconBlockNodeError, PublishOutcome,
+};
+use db::ClientDB;
+use slot_clock::SlotClock;
+use std::sync::{Arc, RwLock};
+use types::{BeaconBlock, PublicKey, Signature};
+
+pub struct BenchingBeaconNode<T: ClientDB, U: SlotClock> {
+    beacon_chain: Arc<BeaconChain<T, U>>,
+    published_blocks: RwLock<Vec<BeaconBlock>>,
+}
+
+impl<T: ClientDB, U: SlotClock> BenchingBeaconNode<T, U> {
+    pub fn new(beacon_chain: Arc<BeaconChain<T, U>>) -> Self {
+        Self {
+            beacon_chain,
+            published_blocks: RwLock::new(vec![]),
+        }
+    }
+
+    pub fn last_published_block(&self) -> Option<BeaconBlock> {
+        Some(
+            self.published_blocks
+                .read()
+                .expect("Unable to unlock `published_blocks` for reading.")
+                .last()?
+                .clone(),
+        )
+    }
+}
+
+impl<T: ClientDB, U: SlotClock> BeaconBlockNode for BenchingBeaconNode<T, U>
+where
+    BlockProductionError: From<<U>::Error>,
+    ProcessingError: From<<U as SlotClock>::Error>,
+{
+    /// Requests the `proposer_nonce` from the `BeaconChain`.
+    fn proposer_nonce(&self, pubkey: &PublicKey) -> Result<u64, BeaconBlockNodeError> {
+        let validator_index = self
+            .beacon_chain
+            .validator_index(pubkey)
+            .ok_or_else(|| BeaconBlockNodeError::RemoteFailure("pubkey unknown.".to_string()))?;
+
+        self.beacon_chain
+            .proposer_slots(validator_index)
+            .ok_or_else(|| {
+                BeaconBlockNodeError::RemoteFailure("validator_index unknown.".to_string())
+            })
+    }
+
+    /// Requests a new `BeaconBlock from the `BeaconChain`.
+    fn produce_beacon_block(
+        &self,
+        slot: u64,
+        randao_reveal: &Signature,
+    ) -> Result<Option<BeaconBlock>, BeaconBlockNodeError>
+where {
+        let (block, _state) = self
+            .beacon_chain
+            .produce_block(randao_reveal.clone())
+            .map_err(|e| BeaconBlockNodeError::RemoteFailure(format!("{:?}", e)))?;
+
+        if block.slot == slot {
+            Ok(Some(block))
+        } else {
+            Err(BeaconBlockNodeError::RemoteFailure(
+                "Unable to produce at non-current slot.".to_string(),
+            ))
+        }
+    }
+
+    /// A block is not _actually_ published to the `BeaconChain`, instead it is stored in the
+    /// `published_block_vec` and a successful `ValidBlock` is returned to the caller.
+    ///
+    /// The block may be retrieved and then applied to the `BeaconChain` manually, potentially in a
+    /// benchmarking scenario.
+    fn publish_beacon_block(
+        &self,
+        block: BeaconBlock,
+    ) -> Result<PublishOutcome, BeaconBlockNodeError> {
+        self.published_blocks
+            .write()
+            .expect("Unable to unlock `published_blocks` for writing.")
+            .push(block);
+        Ok(PublishOutcome::ValidBlock)
+    }
+}

beacon_node/beacon_chain/tests/utils/mod.rs

@@ -1,8 +1,10 @@
+mod benching_beacon_node;
 mod direct_beacon_node;
 mod direct_duties;
 mod test_rig;
 mod validator;
 
+pub use self::benching_beacon_node::BenchingBeaconNode;
 pub use self::direct_beacon_node::DirectBeaconNode;
 pub use self::direct_duties::DirectDuties;
 pub use self::test_rig::TestRig;

beacon_node/beacon_chain/tests/utils/test_rig.rs

@@ -1,6 +1,7 @@
 use super::TestValidator;
 pub use beacon_chain::dump::{Error as DumpError, SlotDump};
 use beacon_chain::BeaconChain;
+use block_producer::BeaconNode;
 #[cfg(test)]
 use db::{
     stores::{BeaconBlockStore, BeaconStateStore},
@@ -11,7 +12,7 @@ use slot_clock::TestingSlotClock;
 use std::fs::File;
 use std::io::prelude::*;
 use std::sync::Arc;
-use types::{ChainSpec, Keypair, Validator};
+use types::{BeaconBlock, ChainSpec, Keypair, Validator};
 
 pub struct TestRig {
     db: Arc<MemoryDB>,
@@ -76,7 +77,12 @@ impl TestRig {
         }
     }
 
-    pub fn produce_next_slot(&mut self) {
+    pub fn advance_chain_with_block(&mut self) {
+        let block = self.produce_next_slot();
+        self.beacon_chain.process_block(block).unwrap();
+    }
+
+    fn produce_next_slot(&mut self) -> BeaconBlock {
         let slot = self
             .beacon_chain
             .present_slot()
@@ -91,7 +97,7 @@ impl TestRig {
             .expect("Unable to determine proposer.");
 
         self.validators[proposer].set_slot(slot);
-        self.validators[proposer].produce_block().unwrap();
+        self.validators[proposer].produce_block().unwrap()
     }
 
     pub fn chain_dump(&self) -> Result<Vec<SlotDump>, DumpError> {
@@ -101,6 +107,7 @@ impl TestRig {
     pub fn dump_to_file(&self, filename: String, chain_dump: &Vec<SlotDump>) {
         let json = serde_json::to_string(chain_dump).unwrap();
         let mut file = File::create(filename).unwrap();
-        file.write_all(json.as_bytes());
+        file.write_all(json.as_bytes())
+            .expect("Failed writing dump to file.");
     }
 }

beacon_node/beacon_chain/tests/utils/validator.rs

@@ -1,11 +1,11 @@
-use super::{DirectBeaconNode, DirectDuties};
+use super::{BenchingBeaconNode, DirectDuties};
 use beacon_chain::BeaconChain;
 #[cfg(test)]
 use block_producer::{test_utils::TestSigner, BlockProducer, Error as PollError};
 use db::MemoryDB;
 use slot_clock::TestingSlotClock;
 use std::sync::Arc;
-use types::{ChainSpec, Keypair};
+use types::{BeaconBlock, ChainSpec, Keypair};
 
 pub use block_producer::PollOutcome;
 
@@ -18,14 +18,14 @@ pub enum ProduceError {
 pub struct TestValidator {
     block_producer: BlockProducer<
         TestingSlotClock,
-        DirectBeaconNode<MemoryDB, TestingSlotClock>,
+        BenchingBeaconNode<MemoryDB, TestingSlotClock>,
         DirectDuties<MemoryDB, TestingSlotClock>,
         TestSigner,
     >,
     spec: Arc<ChainSpec>,
     epoch_map: Arc<DirectDuties<MemoryDB, TestingSlotClock>>,
     keypair: Keypair,
-    beacon_node: Arc<DirectBeaconNode<MemoryDB, TestingSlotClock>>,
+    beacon_node: Arc<BenchingBeaconNode<MemoryDB, TestingSlotClock>>,
     slot_clock: Arc<TestingSlotClock>,
     signer: Arc<TestSigner>,
 }
@@ -38,7 +38,7 @@ impl TestValidator {
         let spec = Arc::new(ChainSpec::foundation());
         let slot_clock = Arc::new(TestingSlotClock::new(0));
         let signer = Arc::new(TestSigner::new(keypair.clone()));
-        let beacon_node = Arc::new(DirectBeaconNode::new(beacon_chain.clone()));
+        let beacon_node = Arc::new(BenchingBeaconNode::new(beacon_chain.clone()));
         let epoch_map = Arc::new(DirectDuties::new(keypair.pk.clone(), beacon_chain.clone()));
 
         let block_producer = BlockProducer::new(
@@ -61,12 +61,18 @@ impl TestValidator {
         }
     }
 
-    pub fn produce_block(&mut self) -> Result<PollOutcome, ProduceError> {
+    pub fn produce_block(&mut self) -> Result<BeaconBlock, ProduceError> {
+        // Using `BenchingBeaconNode`, the validator will always return sucessufully if it tries to
+        // publish a block.
         match self.block_producer.poll() {
-            Ok(PollOutcome::BlockProduced(slot)) => Ok(PollOutcome::BlockProduced(slot)),
+            Ok(PollOutcome::BlockProduced(_)) => {}
-            Ok(outcome) => Err(ProduceError::DidNotProduce(outcome)),
+            Ok(outcome) => return Err(ProduceError::DidNotProduce(outcome)),
-            Err(error) => Err(ProduceError::PollError(error)),
+            Err(error) => return Err(ProduceError::PollError(error)),
-        }
+        };
+        Ok(self
+            .beacon_node
+            .last_published_block()
+            .expect("Unable to obtain produced block."))
     }
 
     pub fn set_slot(&mut self, slot: u64) {