diff --git a/lighthouse/db/src/disk_db.rs b/lighthouse/db/src/disk_db.rs index 2b101355d..b084f483f 100644 --- a/lighthouse/db/src/disk_db.rs +++ b/lighthouse/db/src/disk_db.rs @@ -1,17 +1,10 @@ extern crate rocksdb; +use super::rocksdb::Error as RocksError; +use super::rocksdb::{Options, DB}; +use super::{ClientDB, DBError, DBValue}; use std::fs; use std::path::Path; -use super::rocksdb::{ - DB, - Options, -}; -use super::rocksdb::Error as RocksError; -use super::{ - ClientDB, - DBValue, - DBError -}; /// A on-disk database which implements the ClientDB trait. /// @@ -42,8 +35,7 @@ impl DiskDB { /* * Initialise the path */ - fs::create_dir_all(&path) - .unwrap_or_else(|_| panic!("Unable to create {:?}", &path)); + fs::create_dir_all(&path).unwrap_or_else(|_| panic!("Unable to create {:?}", &path)); let db_path = path.join("database"); /* @@ -51,31 +43,28 @@ impl DiskDB { */ let db = match columns { None => DB::open(&options, db_path), - Some(columns) => DB::open_cf(&options, db_path, columns) + Some(columns) => DB::open_cf(&options, db_path, columns), }.expect("Unable to open local database");; - Self { - db, - } + Self { db } } /// Create a RocksDB column family. Corresponds to the /// `create_cf()` function on the RocksDB API. #[allow(dead_code)] - fn create_col(&mut self, col: &str) - -> Result<(), DBError> - { + fn create_col(&mut self, col: &str) -> Result<(), DBError> { match self.db.create_cf(col, &Options::default()) { Err(e) => Err(e.into()), - Ok(_) => Ok(()) + Ok(_) => Ok(()), } } - } impl From for DBError { fn from(e: RocksError) -> Self { - Self { message: e.to_string() } + Self { + message: e.to_string(), + } } } @@ -85,17 +74,15 @@ impl ClientDB for DiskDB { /// Corresponds to the `get_cf()` method on the RocksDB API. /// Will attempt to get the `ColumnFamily` and return an Err /// if it fails. - fn get(&self, col: &str, key: &[u8]) - -> Result, DBError> - { + fn get(&self, col: &str, key: &[u8]) -> Result, DBError> { match self.db.cf_handle(col) { - None => Err(DBError{ message: "Unknown column".to_string() }), - Some(handle) => { - match self.db.get_cf(handle, key)? { - None => Ok(None), - Some(db_vec) => Ok(Some(DBValue::from(&*db_vec))) - } - } + None => Err(DBError { + message: "Unknown column".to_string(), + }), + Some(handle) => match self.db.get_cf(handle, key)? { + None => Ok(None), + Some(db_vec) => Ok(Some(DBValue::from(&*db_vec))), + }, } } @@ -104,38 +91,54 @@ impl ClientDB for DiskDB { /// Corresponds to the `cf_handle()` method on the RocksDB API. /// Will attempt to get the `ColumnFamily` and return an Err /// if it fails. - fn put(&self, col: &str, key: &[u8], val: &[u8]) - -> Result<(), DBError> - { + fn put(&self, col: &str, key: &[u8], val: &[u8]) -> Result<(), DBError> { match self.db.cf_handle(col) { - None => Err(DBError{ message: "Unknown column".to_string() }), - Some(handle) => self.db.put_cf(handle, key, val).map_err(|e| e.into()) + None => Err(DBError { + message: "Unknown column".to_string(), + }), + Some(handle) => self.db.put_cf(handle, key, val).map_err(|e| e.into()), } } /// Return true if some key exists in some column. - fn exists(&self, col: &str, key: &[u8]) - -> Result - { + fn exists(&self, col: &str, key: &[u8]) -> Result { /* * I'm not sure if this is the correct way to read if some - * block exists. Naievely I would expect this to unncessarily + * block exists. Naively I would expect this to unncessarily * copy some data, but I could be wrong. */ match self.db.cf_handle(col) { - None => Err(DBError{ message: "Unknown column".to_string() }), - Some(handle) => Ok(self.db.get_cf(handle, key)?.is_some()) + None => Err(DBError { + message: "Unknown column".to_string(), + }), + Some(handle) => Ok(self.db.get_cf(handle, key)?.is_some()), + } + } + + /// Delete the value for some key on some column. + /// + /// Corresponds to the `delete_cf()` method on the RocksDB API. + /// Will attempt to get the `ColumnFamily` and return an Err + /// if it fails. + fn delete(&self, col: &str, key: &[u8]) -> Result<(), DBError> { + match self.db.cf_handle(col) { + None => Err(DBError { + message: "Unknown column".to_string(), + }), + Some(handle) => { + self.db.delete_cf(handle, key)?; + Ok(()) + } } } } - #[cfg(test)] mod tests { - use super::*; use super::super::ClientDB; - use std::{ env, fs, thread }; + use super::*; use std::sync::Arc; + use std::{env, fs, thread}; #[test] #[ignore] diff --git a/lighthouse/db/src/memory_db.rs b/lighthouse/db/src/memory_db.rs index c912b7c68..008e5912f 100644 --- a/lighthouse/db/src/memory_db.rs +++ b/lighthouse/db/src/memory_db.rs @@ -1,12 +1,8 @@ -use std::collections::{ HashSet, HashMap }; -use std::sync::RwLock; use super::blake2::blake2b::blake2b; use super::COLUMNS; -use super::{ - ClientDB, - DBValue, - DBError -}; +use super::{ClientDB, DBError, DBValue}; +use std::collections::{HashMap, HashSet}; +use std::sync::RwLock; type DBHashMap = HashMap, Vec>; type ColumnHashSet = HashSet; @@ -17,7 +13,7 @@ type ColumnHashSet = HashSet; /// this DB would be used outside of tests. pub struct MemoryDB { db: RwLock, - known_columns: RwLock + known_columns: RwLock, } impl MemoryDB { @@ -45,9 +41,7 @@ impl MemoryDB { impl ClientDB for MemoryDB { /// Get the value of some key from the database. Returns `None` if the key does not exist. - fn get(&self, col: &str, key: &[u8]) - -> Result, DBError> - { + fn get(&self, col: &str, key: &[u8]) -> Result, DBError> { // Panic if the DB locks are poisoned. let db = self.db.read().unwrap(); let known_columns = self.known_columns.read().unwrap(); @@ -56,14 +50,14 @@ impl ClientDB for MemoryDB { let column_key = MemoryDB::get_key_for_col(col, key); Ok(db.get(&column_key).and_then(|val| Some(val.clone()))) } else { - Err(DBError{ message: "Unknown column".to_string() }) + Err(DBError { + message: "Unknown column".to_string(), + }) } } /// Puts a key in the database. - fn put(&self, col: &str, key: &[u8], val: &[u8]) - -> Result<(), DBError> - { + fn put(&self, col: &str, key: &[u8], val: &[u8]) -> Result<(), DBError> { // Panic if the DB locks are poisoned. let mut db = self.db.write().unwrap(); let known_columns = self.known_columns.read().unwrap(); @@ -73,14 +67,14 @@ impl ClientDB for MemoryDB { db.insert(column_key, val.to_vec()); Ok(()) } else { - Err(DBError{ message: "Unknown column".to_string() }) + Err(DBError { + message: "Unknown column".to_string(), + }) } } /// Return true if some key exists in some column. - fn exists(&self, col: &str, key: &[u8]) - -> Result - { + fn exists(&self, col: &str, key: &[u8]) -> Result { // Panic if the DB locks are poisoned. let db = self.db.read().unwrap(); let known_columns = self.known_columns.read().unwrap(); @@ -89,23 +83,55 @@ impl ClientDB for MemoryDB { let column_key = MemoryDB::get_key_for_col(col, key); Ok(db.contains_key(&column_key)) } else { - Err(DBError{ message: "Unknown column".to_string() }) + Err(DBError { + message: "Unknown column".to_string(), + }) + } + } + /// Delete some key from the database. + fn delete(&self, col: &str, key: &[u8]) -> Result<(), DBError> { + // Panic if the DB locks are poisoned. + let mut db = self.db.write().unwrap(); + let known_columns = self.known_columns.read().unwrap(); + + if known_columns.contains(&col.to_string()) { + let column_key = MemoryDB::get_key_for_col(col, key); + db.remove(&column_key); + Ok(()) + } else { + Err(DBError { + message: "Unknown column".to_string(), + }) } } } - #[cfg(test)] mod tests { - use super::*; + use super::super::stores::{BLOCKS_DB_COLUMN, VALIDATOR_DB_COLUMN}; use super::super::ClientDB; - use std::thread; + use super::*; use std::sync::Arc; - use super::super::stores::{ - BLOCKS_DB_COLUMN, - VALIDATOR_DB_COLUMN, - }; + use std::thread; + + #[test] + fn test_memorydb_can_delete() { + let col_a: &str = BLOCKS_DB_COLUMN; + + let db = MemoryDB::open(); + + db.put(col_a, "dogs".as_bytes(), "lol".as_bytes()).unwrap(); + + assert_eq!( + db.get(col_a, "dogs".as_bytes()).unwrap().unwrap(), + "lol".as_bytes() + ); + + db.delete(col_a, "dogs".as_bytes()).unwrap(); + + assert_eq!(db.get(col_a, "dogs".as_bytes()).unwrap(), None); + } #[test] fn test_memorydb_column_access() { @@ -121,10 +147,14 @@ mod tests { db.put(col_a, "same".as_bytes(), "cat".as_bytes()).unwrap(); db.put(col_b, "same".as_bytes(), "dog".as_bytes()).unwrap(); - assert_eq!(db.get(col_a, "same".as_bytes()).unwrap().unwrap(), "cat".as_bytes()); - assert_eq!(db.get(col_b, "same".as_bytes()).unwrap().unwrap(), "dog".as_bytes()); - - + assert_eq!( + db.get(col_a, "same".as_bytes()).unwrap().unwrap(), + "cat".as_bytes() + ); + assert_eq!( + db.get(col_b, "same".as_bytes()).unwrap().unwrap(), + "dog".as_bytes() + ); } #[test] diff --git a/lighthouse/db/src/traits.rs b/lighthouse/db/src/traits.rs index 25eaf3abe..41be3e23d 100644 --- a/lighthouse/db/src/traits.rs +++ b/lighthouse/db/src/traits.rs @@ -2,7 +2,7 @@ pub type DBValue = Vec; #[derive(Debug)] pub struct DBError { - pub message: String + pub message: String, } impl DBError { @@ -18,13 +18,11 @@ impl DBError { /// program to use a persistent on-disk database during production, /// but use a transient database during tests. pub trait ClientDB: Sync + Send { - fn get(&self, col: &str, key: &[u8]) - -> Result, DBError>; + fn get(&self, col: &str, key: &[u8]) -> Result, DBError>; - fn put(&self, col: &str, key: &[u8], val: &[u8]) - -> Result<(), DBError>; + fn put(&self, col: &str, key: &[u8], val: &[u8]) -> Result<(), DBError>; - fn exists(&self, col: &str, key: &[u8]) - -> Result; + fn exists(&self, col: &str, key: &[u8]) -> Result; + + fn delete(&self, col: &str, key: &[u8]) -> Result<(), DBError>; } -