lighthouse/beacon_node/network/src/persisted_dht.rs
Adam Szkoda 9db0c28051
Make key value storage abstractions more accurate (#1267)
* Layer do_atomically() abstractions properly

* Reduce allocs and DRY get_key_for_col()

* Parameterize HotColdDB with hot and cold item stores

* -impl Store for MemoryStore

* Replace Store uses with HotColdDB

* Ditch Store trait

* cargo fmt

* Style fix

* Readd missing dep that broke the build
2020-06-16 11:34:04 +10:00

83 lines
2.6 KiB
Rust

use eth2_libp2p::Enr;
use rlp;
use std::sync::Arc;
use store::{DBColumn, Error as StoreError, HotColdDB, ItemStore, StoreItem};
use types::{EthSpec, Hash256};
/// 32-byte key for accessing the `DhtEnrs`.
pub const DHT_DB_KEY: &str = "PERSISTEDDHTPERSISTEDDHTPERSISTE";
pub fn load_dht<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
store: Arc<HotColdDB<E, Hot, Cold>>,
) -> Vec<Enr> {
// Load DHT from store
let key = Hash256::from_slice(&DHT_DB_KEY.as_bytes());
match store.get_item(&key) {
Ok(Some(p)) => {
let p: PersistedDht = p;
p.enrs
}
_ => Vec::new(),
}
}
/// Attempt to persist the ENR's in the DHT to `self.store`.
pub fn persist_dht<E: EthSpec, Hot: ItemStore<E>, Cold: ItemStore<E>>(
store: Arc<HotColdDB<E, Hot, Cold>>,
enrs: Vec<Enr>,
) -> Result<(), store::Error> {
let key = Hash256::from_slice(&DHT_DB_KEY.as_bytes());
store.put_item(&key, &PersistedDht { enrs })?;
Ok(())
}
/// Wrapper around DHT for persistence to disk.
pub struct PersistedDht {
pub enrs: Vec<Enr>,
}
impl StoreItem for PersistedDht {
fn db_column() -> DBColumn {
DBColumn::DhtEnrs
}
fn as_store_bytes(&self) -> Vec<u8> {
rlp::encode_list(&self.enrs)
}
fn from_store_bytes(bytes: &[u8]) -> Result<Self, StoreError> {
let rlp = rlp::Rlp::new(bytes);
let enrs: Vec<Enr> = rlp
.as_list()
.map_err(|e| StoreError::RlpError(format!("{}", e)))?;
Ok(PersistedDht { enrs })
}
}
#[cfg(test)]
mod tests {
use super::*;
use eth2_libp2p::Enr;
use sloggers::{null::NullLoggerBuilder, Build};
use std::str::FromStr;
use store::config::StoreConfig;
use store::{HotColdDB, MemoryStore};
use types::{ChainSpec, Hash256, MinimalEthSpec};
#[test]
fn test_persisted_dht() {
let log = NullLoggerBuilder.build().unwrap();
let store: HotColdDB<
MinimalEthSpec,
MemoryStore<MinimalEthSpec>,
MemoryStore<MinimalEthSpec>,
> = HotColdDB::open_ephemeral(StoreConfig::default(), ChainSpec::minimal(), log).unwrap();
let enrs = vec![Enr::from_str("enr:-IS4QHCYrYZbAKWCBRlAy5zzaDZXJBGkcnh4MHcBFZntXNFrdvJjX04jRzjzCBOonrkTfj499SZuOh8R33Ls8RRcy5wBgmlkgnY0gmlwhH8AAAGJc2VjcDI1NmsxoQPKY0yuDUmstAHYpMa2_oxVtw0RW_QAdpzBQA8yWM0xOIN1ZHCCdl8").unwrap()];
let key = Hash256::from_slice(&DHT_DB_KEY.as_bytes());
store
.put_item(&key, &PersistedDht { enrs: enrs.clone() })
.unwrap();
let dht: PersistedDht = store.get_item(&key).unwrap().unwrap();
assert_eq!(dht.enrs, enrs);
}
}