Expand beacon_chain_sim

tests/beacon_chain_sim/src/main.rs

@@ -1,19 +1,23 @@
 use node_test_rig::{
     environment::{EnvironmentBuilder, RuntimeContext},
-    testing_client_config, ClientConfig, LocalBeaconNode, ProductionClient,
+    testing_client_config, ClientConfig, LocalBeaconNode, LocalValidatorClient, ProductionClient,
+    ValidatorConfig,
 };
 use types::EthSpec;
 
 pub type BeaconNode<E> = LocalBeaconNode<ProductionClient<E>>;
 
 fn main() {
-    match simulation(4) {
+    let nodes = 4;
+    let validators_per_node = 64 / nodes;
+
+    match simulation(nodes, validators_per_node) {
         Ok(()) => println!("Simulation exited successfully"),
         Err(e) => println!("Simulation exited with error: {}", e),
     }
 }
 
-fn simulation(num_nodes: usize) -> Result<(), String> {
+fn simulation(num_nodes: usize, validators_per_node: usize) -> Result<(), String> {
     if num_nodes < 1 {
         return Err("Must have at least one node".into());
     }
@@ -28,19 +32,48 @@ fn simulation(num_nodes: usize) -> Result<(), String> {
     let boot_node =
         BeaconNode::production(env.service_context("boot_node".into()), base_config.clone());
 
-    let nodes = (1..num_nodes)
+    let mut nodes = (1..num_nodes)
         .map(|i| {
             let context = env.service_context(format!("node_{}", i));
             new_with_bootnode_via_enr(context, &boot_node, base_config.clone())
         })
         .collect::<Vec<_>>();
 
+    let validators = nodes
+        .iter()
+        .enumerate()
+        .map(|(i, node)| {
+            let mut context = env.service_context(format!("validator_{}", i));
+
+            // Pull the spec from the beacon node's beacon chain, in case there were some changes
+            // to the spec after the node booted.
+            context.eth2_config.spec = node
+                .client
+                .beacon_chain()
+                .expect("should have beacon chain")
+                .spec
+                .clone();
+
+            let indices =
+                (i * validators_per_node..(i + 1) * validators_per_node).collect::<Vec<_>>();
+            new_validator_client(
+                env.service_context(format!("validator_{}", i)),
+                node,
+                ValidatorConfig::default(),
+                &indices,
+            )
+        })
+        .collect::<Vec<_>>();
+
+    nodes.insert(0, boot_node);
+
     env.block_until_ctrl_c()?;
 
     Ok(())
 }
 
-// TODO: this function does not result in nodes connecting to each other. Age to investigate?
+// TODO: this function does not result in nodes connecting to each other. This is a bug due to
+// using a 0 port for discovery. Age is fixing it.
 fn new_with_bootnode_via_enr<E: EthSpec>(
     context: RuntimeContext<E>,
     boot_node: &BeaconNode<E>,
@@ -75,3 +108,22 @@ fn new_with_bootnode_via_multiaddr<E: EthSpec>(
 
     BeaconNode::production(context, config)
 }
+
+fn new_validator_client<E: EthSpec>(
+    context: RuntimeContext<E>,
+    beacon_node: &BeaconNode<E>,
+    base_config: ValidatorConfig,
+    keypair_indices: &[usize],
+) -> LocalValidatorClient<E> {
+    let mut config = base_config;
+
+    let (grpc_endpoint, grpc_port) = beacon_node
+        .client
+        .grpc_listen_addr()
+        .expect("Must have gRPC started");
+
+    config.server = grpc_endpoint;
+    config.server_grpc_port = grpc_port;
+
+    LocalValidatorClient::production_with_insecure_keypairs(context, config, keypair_indices)
+}