Modify runtime to allow memory or disk db

DiskDB is not working yet, but we'll get there!

beacon_node/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2018"
 
 [dependencies]
 types = { path = "../eth2/types" }
+db = { path = "./db" }
 client = { path = "client" }
 version = { path = "version" }
 clap = "2.32.0"

beacon_node/beacon_chain/src/initialise.rs

@@ -61,7 +61,7 @@ pub fn initialise_beacon_chain(
 }
 
 /// Initialisation of a test beacon chain, uses an in memory db with fixed genesis time.
-pub fn initialise_test_beacon_chain(
+pub fn initialise_test_beacon_chain_with_memory_db(
     spec: &ChainSpec,
     _db_name: Option<&PathBuf>,
 ) -> Arc<BeaconChain<MemoryDB, SystemTimeSlotClock, BitwiseLMDGhost<MemoryDB>>> {
@@ -100,3 +100,44 @@ pub fn initialise_test_beacon_chain(
         .expect("Terminate if beacon chain generation fails"),
     )
 }
+
+/// Initialisation of a test beacon chain, uses an in memory db with fixed genesis time.
+pub fn initialise_test_beacon_chain_with_disk_db(
+    spec: &ChainSpec,
+    db_name: Option<&PathBuf>,
+) -> Arc<BeaconChain<DiskDB, SystemTimeSlotClock, BitwiseLMDGhost<DiskDB>>> {
+    let db = Arc::new(DiskDB::open(db_name.expect("Must have DB path"), None));
+    let block_store = Arc::new(BeaconBlockStore::new(db.clone()));
+    let state_store = Arc::new(BeaconStateStore::new(db.clone()));
+
+    let state_builder = TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(8, spec);
+    let (genesis_state, _keypairs) = state_builder.build();
+
+    let mut genesis_block = BeaconBlock::empty(spec);
+    genesis_block.state_root = Hash256::from_slice(&genesis_state.hash_tree_root());
+
+    // Slot clock
+    let slot_clock = SystemTimeSlotClock::new(
+        spec.genesis_slot,
+        genesis_state.genesis_time,
+        spec.seconds_per_slot,
+    )
+    .expect("Unable to load SystemTimeSlotClock");
+    // Choose the fork choice
+    let fork_choice = BitwiseLMDGhost::new(block_store.clone(), state_store.clone());
+
+    // Genesis chain
+    //TODO: Handle error correctly
+    Arc::new(
+        BeaconChain::from_genesis(
+            state_store.clone(),
+            block_store.clone(),
+            slot_clock,
+            genesis_state,
+            genesis_block,
+            spec.clone(),
+            fork_choice,
+        )
+        .expect("Terminate if beacon chain generation fails"),
+    )
+}

beacon_node/client/src/client_config.rs

@@ -119,6 +119,12 @@ impl ClientConfig {
             }
         }
 
+        match args.value_of("db") {
+            Some("rocks") => config.db_type = DBType::RocksDB,
+            Some("memory") => config.db_type = DBType::Memory,
+            _ => unreachable!(), // clap prevents this.
+        };
+
         Ok(config)
     }
 }

beacon_node/client/src/client_types.rs

@@ -34,9 +34,9 @@ impl ClientTypes for StandardClientType {
     }
 }
 
-pub struct TestingClientType;
+pub struct MemoryDBTestingClientType;
 
-impl ClientTypes for TestingClientType {
+impl ClientTypes for MemoryDBTestingClientType {
     type DB = MemoryDB;
     type SlotClock = SystemTimeSlotClock;
     type ForkChoice = BitwiseLMDGhost<MemoryDB>;
@@ -44,6 +44,20 @@ impl ClientTypes for TestingClientType {
     fn initialise_beacon_chain(
         config: &ClientConfig,
     ) -> Arc<BeaconChain<Self::DB, Self::SlotClock, Self::ForkChoice>> {
-        initialise::initialise_test_beacon_chain(&config.spec, None)
+        initialise::initialise_test_beacon_chain_with_memory_db(&config.spec, None)
+    }
+}
+
+pub struct DiskDBTestingClientType;
+
+impl ClientTypes for DiskDBTestingClientType {
+    type DB = DiskDB;
+    type SlotClock = SystemTimeSlotClock;
+    type ForkChoice = BitwiseLMDGhost<DiskDB>;
+
+    fn initialise_beacon_chain(
+        config: &ClientConfig,
+    ) -> Arc<BeaconChain<Self::DB, Self::SlotClock, Self::ForkChoice>> {
+        initialise::initialise_test_beacon_chain_with_disk_db(&config.spec, Some(&config.db_name))
     }
 }

beacon_node/src/main.rs

@@ -58,6 +58,15 @@ fn main() {
                 .help("Listen port for RPC endpoint.")
                 .takes_value(true),
         )
+        .arg(
+            Arg::with_name("db")
+                .long("db")
+                .value_name("DB")
+                .help("Type of database to use.")
+                .takes_value(true)
+                .possible_values(&["rocks", "memory"])
+                .default_value("memory"),
+        )
         .get_matches();
 
     // invalid arguments, panic

beacon_node/src/run.rs

@@ -1,15 +1,18 @@
-use client::client_types::TestingClientType;
+use client::client_types::{DiskDBTestingClientType, MemoryDBTestingClientType};
 use client::error;
-use client::{notifier, Client, ClientConfig};
+use client::{notifier, Client, ClientConfig, ClientTypes};
+use db::DBType;
 use futures::sync::oneshot;
 use futures::Future;
 use slog::info;
 use std::cell::RefCell;
 use tokio::runtime::Builder;
+use tokio::runtime::Runtime;
+use tokio::runtime::TaskExecutor;
 use tokio_timer::clock::Clock;
 
 pub fn run_beacon_node(config: ClientConfig, log: &slog::Logger) -> error::Result<()> {
-    let mut runtime = Builder::new()
+    let runtime = Builder::new()
         .name_prefix("main-")
         .clock(Clock::system())
         .build()
@@ -20,8 +23,32 @@ pub fn run_beacon_node(config: ClientConfig, log: &slog::Logger) -> error::Resul
           "data_dir" => &config.data_dir.to_str(),
           "port" => &config.net_conf.listen_port);
 
+    let executor = runtime.executor();
+
+    match config.db_type {
+        DBType::RocksDB => {
+            let client: Client<DiskDBTestingClientType> =
+                Client::new(config, log.clone(), &executor)?;
+
+            run(client, executor, runtime, log)
+        }
+        DBType::Memory => {
+            let client: Client<MemoryDBTestingClientType> =
+                Client::new(config, log.clone(), &executor)?;
+
+            run(client, executor, runtime, log)
+        }
+    }
+}
+
+pub fn run<T: ClientTypes>(
+    client: Client<T>,
+    executor: TaskExecutor,
+    mut runtime: Runtime,
+    log: &slog::Logger,
+) -> error::Result<()> {
     // run service until ctrl-c
-    let (ctrlc_send, ctrlc) = oneshot::channel();
+    let (ctrlc_send, ctrlc_oneshot) = oneshot::channel();
     let ctrlc_send_c = RefCell::new(Some(ctrlc_send));
     ctrlc::set_handler(move || {
         if let Some(ctrlc_send) = ctrlc_send_c.try_borrow_mut().unwrap().take() {
@@ -32,14 +59,10 @@ pub fn run_beacon_node(config: ClientConfig, log: &slog::Logger) -> error::Resul
 
     let (exit_signal, exit) = exit_future::signal();
 
-    let executor = runtime.executor();
-
-    // currently testing - using TestingClientType
-    let client: Client<TestingClientType> = Client::new(config, log.clone(), &executor)?;
     notifier::run(&client, executor, exit);
 
     runtime
-        .block_on(ctrlc)
+        .block_on(ctrlc_oneshot)
         .map_err(|e| format!("Ctrlc oneshot failed: {:?}", e))?;
 
     // perform global shutdown operations.