mod test;
use crate::K_VALUE;
use crate::addresses::Addresses;
use crate::handler::{KademliaHandler, KademliaRequestId, KademliaHandlerEvent, KademliaHandlerIn};
use crate::jobs::*;
use crate::kbucket::{self, KBucketsTable, NodeStatus};
use crate::protocol::{KadConnectionType, KadPeer};
use crate::query::{Query, QueryId, QueryPool, QueryConfig, QueryPoolState};
use crate::record::{self, store::{self, RecordStore}, Record, ProviderRecord};
use fnv::{FnvHashMap, FnvHashSet};
use libp2p_core::{ConnectedPoint, Multiaddr, PeerId};
use libp2p_swarm::{NetworkBehaviour, NetworkBehaviourAction, PollParameters, ProtocolsHandler};
use log::{info, debug, warn};
use smallvec::SmallVec;
use std::{borrow::{Borrow, Cow}, error, iter, time::Duration};
use std::collections::VecDeque;
use std::num::NonZeroUsize;
use std::task::{Context, Poll};
use wasm_timer::Instant;
pub struct Kademlia<TStore> {
kbuckets: KBucketsTable<kbucket::Key<PeerId>, Addresses>,
protocol_name_override: Option<Cow<'static, [u8]>>,
queries: QueryPool<QueryInner>,
connected_peers: FnvHashSet<PeerId>,
add_provider_job: Option<AddProviderJob>,
put_record_job: Option<PutRecordJob>,
record_ttl: Option<Duration>,
provider_record_ttl: Option<Duration>,
queued_events: VecDeque<NetworkBehaviourAction<KademliaHandlerIn<QueryId>, KademliaEvent>>,
store: TStore,
}
#[derive(Debug, Clone)]
pub struct KademliaConfig {
kbucket_pending_timeout: Duration,
query_config: QueryConfig,
protocol_name_override: Option<Cow<'static, [u8]>>,
record_ttl: Option<Duration>,
record_replication_interval: Option<Duration>,
record_publication_interval: Option<Duration>,
provider_record_ttl: Option<Duration>,
provider_publication_interval: Option<Duration>,
}
impl Default for KademliaConfig {
fn default() -> Self {
KademliaConfig {
kbucket_pending_timeout: Duration::from_secs(60),
query_config: QueryConfig::default(),
protocol_name_override: None,
record_ttl: Some(Duration::from_secs(36 * 60 * 60)),
record_replication_interval: Some(Duration::from_secs(60 * 60)),
record_publication_interval: Some(Duration::from_secs(24 * 60 * 60)),
provider_publication_interval: Some(Duration::from_secs(12 * 60 * 60)),
provider_record_ttl: Some(Duration::from_secs(24 * 60 * 60)),
}
}
}
impl KademliaConfig {
pub fn set_protocol_name(&mut self, name: impl Into<Cow<'static, [u8]>>) -> &mut Self {
self.protocol_name_override = Some(name.into());
self
}
pub fn set_query_timeout(&mut self, timeout: Duration) -> &mut Self {
self.query_config.timeout = timeout;
self
}
pub fn set_replication_factor(&mut self, replication_factor: NonZeroUsize) -> &mut Self {
self.query_config.replication_factor = replication_factor;
self
}
pub fn set_record_ttl(&mut self, record_ttl: Option<Duration>) -> &mut Self {
self.record_ttl = record_ttl;
self
}
pub fn set_replication_interval(&mut self, interval: Option<Duration>) -> &mut Self {
self.record_replication_interval = interval;
self
}
pub fn set_publication_interval(&mut self, interval: Option<Duration>) -> &mut Self {
self.record_publication_interval = interval;
self
}
pub fn set_provider_record_ttl(&mut self, ttl: Option<Duration>) -> &mut Self {
self.provider_record_ttl = ttl;
self
}
pub fn set_provider_publication_interval(&mut self, interval: Option<Duration>) -> &mut Self {
self.provider_publication_interval = interval;
self
}
}
impl<TStore> Kademlia<TStore>
where
for<'a> TStore: RecordStore<'a>
{
pub fn new(id: PeerId, store: TStore) -> Self {
Self::with_config(id, store, Default::default())
}
pub fn with_config(id: PeerId, store: TStore, config: KademliaConfig) -> Self {
let local_key = kbucket::Key::new(id.clone());
let put_record_job = config
.record_replication_interval
.or(config.record_publication_interval)
.map(|interval| PutRecordJob::new(
id.clone(),
interval,
config.record_publication_interval,
config.record_ttl,
));
let add_provider_job = config
.provider_publication_interval
.map(AddProviderJob::new);
Kademlia {
store,
kbuckets: KBucketsTable::new(local_key, config.kbucket_pending_timeout),
protocol_name_override: config.protocol_name_override,
queued_events: VecDeque::with_capacity(config.query_config.replication_factor.get()),
queries: QueryPool::new(config.query_config),
connected_peers: Default::default(),
add_provider_job,
put_record_job,
record_ttl: config.record_ttl,
provider_record_ttl: config.provider_record_ttl,
}
}
pub fn add_address(&mut self, peer: &PeerId, address: Multiaddr) {
let key = kbucket::Key::new(peer.clone());
match self.kbuckets.entry(&key) {
kbucket::Entry::Present(mut entry, _) => {
if entry.value().insert(address) {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::RoutingUpdated {
peer: peer.clone(),
addresses: entry.value().clone(),
old_peer: None,
}
))
}
}
kbucket::Entry::Pending(mut entry, _) => {
entry.value().insert(address);
}
kbucket::Entry::Absent(entry) => {
let addresses = Addresses::new(address);
let status =
if self.connected_peers.contains(peer) {
NodeStatus::Connected
} else {
NodeStatus::Disconnected
};
match entry.insert(addresses.clone(), status) {
kbucket::InsertResult::Inserted => {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::RoutingUpdated {
peer: peer.clone(),
addresses,
old_peer: None,
}
));
},
kbucket::InsertResult::Full => {
debug!("Bucket full. Peer not added to routing table: {}", peer)
},
kbucket::InsertResult::Pending { disconnected } => {
self.queued_events.push_back(NetworkBehaviourAction::DialPeer {
peer_id: disconnected.into_preimage(),
})
},
}
},
kbucket::Entry::SelfEntry => {},
}
}
pub fn kbuckets_entries(&mut self) -> impl Iterator<Item = &PeerId> {
self.kbuckets.iter().map(|entry| entry.node.key.preimage())
}
pub fn get_closest_peers<K>(&mut self, key: K)
where
K: Borrow<[u8]> + Clone
{
let info = QueryInfo::GetClosestPeers { key: key.borrow().to_vec() };
let target = kbucket::Key::new(key);
let peers = self.kbuckets.closest_keys(&target);
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
pub fn get_record(&mut self, key: &record::Key, quorum: Quorum) {
let quorum = quorum.eval(self.queries.config().replication_factor);
let mut records = Vec::with_capacity(quorum.get());
if let Some(record) = self.store.get(key) {
if record.is_expired(Instant::now()) {
self.store.remove(key)
} else {
records.push(record.into_owned());
if quorum.get() == 1 {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::GetRecordResult(Ok(GetRecordOk { records }))
));
return;
}
}
}
let target = kbucket::Key::new(key.clone());
let info = QueryInfo::GetRecord { key: key.clone(), records, quorum, cache_at: None };
let peers = self.kbuckets.closest_keys(&target);
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
pub fn put_record(&mut self, mut record: Record, quorum: Quorum) {
record.publisher = Some(self.kbuckets.local_key().preimage().clone());
if let Err(err) = self.store.put(record.clone()) {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::PutRecordResult(Err(
PutRecordError::LocalStorageError {
key: record.key,
cause: err,
}
))
));
} else {
record.expires = record.expires.or_else(||
self.record_ttl.map(|ttl| Instant::now() + ttl));
let quorum = quorum.eval(self.queries.config().replication_factor);
let target = kbucket::Key::new(record.key.clone());
let peers = self.kbuckets.closest_keys(&target);
let context = PutRecordContext::Publish;
let info = QueryInfo::PreparePutRecord { record, quorum, context };
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
}
pub fn remove_record(&mut self, key: &record::Key) {
if let Some(r) = self.store.get(key) {
if r.publisher.as_ref() == Some(self.kbuckets.local_key().preimage()) {
self.store.remove(key)
}
}
}
pub fn store_mut(&mut self) -> &mut TStore {
&mut self.store
}
pub fn bootstrap(&mut self) {
let local_key = self.kbuckets.local_key().clone();
let info = QueryInfo::Bootstrap { peer: local_key.preimage().clone() };
let peers = self.kbuckets.closest_keys(&local_key).collect::<Vec<_>>();
let inner = QueryInner::new(info);
self.queries.add_iter_closest(local_key, peers, inner);
}
pub fn start_providing(&mut self, key: record::Key) {
let record = ProviderRecord::new(key.clone(), self.kbuckets.local_key().preimage().clone());
if let Err(err) = self.store.add_provider(record) {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::StartProvidingResult(Err(
AddProviderError::LocalStorageError { key, cause: err }
))
));
} else {
let target = kbucket::Key::new(key.clone());
let peers = self.kbuckets.closest_keys(&target);
let context = AddProviderContext::Publish;
let info = QueryInfo::PrepareAddProvider { key, context };
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
}
pub fn stop_providing(&mut self, key: &record::Key) {
self.store.remove_provider(key, self.kbuckets.local_key().preimage());
}
pub fn get_providers(&mut self, key: record::Key) {
let info = QueryInfo::GetProviders {
key: key.clone(),
providers: Vec::new(),
};
let target = kbucket::Key::new(key);
let peers = self.kbuckets.closest_keys(&target);
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
fn discovered<'a, I>(&'a mut self, query_id: &QueryId, source: &PeerId, peers: I)
where
I: Iterator<Item = &'a KadPeer> + Clone
{
let local_id = self.kbuckets.local_key().preimage().clone();
let others_iter = peers.filter(|p| p.node_id != local_id);
for peer in others_iter.clone() {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::Discovered {
peer_id: peer.node_id.clone(),
addresses: peer.multiaddrs.clone(),
ty: peer.connection_ty,
}
));
}
if let Some(query) = self.queries.get_mut(query_id) {
log::trace!("Request to {:?} in query {:?} succeeded.", source, query_id);
for peer in others_iter.clone() {
log::trace!("Peer {:?} reported by {:?} in query {:?}.",
peer, source, query_id);
let addrs = peer.multiaddrs.iter().cloned().collect();
query.inner.addresses.insert(peer.node_id.clone(), addrs);
}
query.on_success(source, others_iter.cloned().map(|kp| kp.node_id))
}
}
fn find_closest<T: Clone>(&mut self, target: &kbucket::Key<T>, source: &PeerId) -> Vec<KadPeer> {
if target == self.kbuckets.local_key() {
Vec::new()
} else {
self.kbuckets
.closest(target)
.filter(|e| e.node.key.preimage() != source)
.take(self.queries.config().replication_factor.get())
.map(KadPeer::from)
.collect()
}
}
fn provider_peers(&mut self, key: &record::Key, source: &PeerId) -> Vec<KadPeer> {
let kbuckets = &mut self.kbuckets;
self.store.providers(key)
.into_iter()
.filter_map(move |p|
if &p.provider != source {
let key = kbucket::Key::new(p.provider.clone());
kbuckets.entry(&key).view().map(|e| KadPeer::from(e.to_owned()))
} else {
None
})
.take(self.queries.config().replication_factor.get())
.collect()
}
fn start_add_provider(&mut self, key: record::Key, context: AddProviderContext) {
let info = QueryInfo::PrepareAddProvider { key: key.clone(), context };
let target = kbucket::Key::new(key);
let peers = self.kbuckets.closest_keys(&target);
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
fn start_put_record(&mut self, record: Record, quorum: Quorum, context: PutRecordContext) {
let quorum = quorum.eval(self.queries.config().replication_factor);
let target = kbucket::Key::new(record.key.clone());
let peers = self.kbuckets.closest_keys(&target);
let info = QueryInfo::PreparePutRecord { record, quorum, context };
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
fn connection_updated(&mut self, peer: PeerId, address: Option<Multiaddr>, new_status: NodeStatus) {
let key = kbucket::Key::new(peer.clone());
match self.kbuckets.entry(&key) {
kbucket::Entry::Present(mut entry, old_status) => {
if let Some(address) = address {
if entry.value().insert(address) {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::RoutingUpdated {
peer,
addresses: entry.value().clone(),
old_peer: None,
}
))
}
}
if old_status != new_status {
entry.update(new_status);
}
},
kbucket::Entry::Pending(mut entry, old_status) => {
if let Some(address) = address {
entry.value().insert(address);
}
if old_status != new_status {
entry.update(new_status);
}
},
kbucket::Entry::Absent(entry) => {
if new_status == NodeStatus::Connected {
if let Some(address) = address {
let addresses = Addresses::new(address);
match entry.insert(addresses.clone(), new_status) {
kbucket::InsertResult::Inserted => {
let event = KademliaEvent::RoutingUpdated {
peer: peer.clone(),
addresses,
old_peer: None,
};
self.queued_events.push_back(
NetworkBehaviourAction::GenerateEvent(event));
},
kbucket::InsertResult::Full => {
debug!("Bucket full. Peer not added to routing table: {}", peer)
},
kbucket::InsertResult::Pending { disconnected } => {
debug_assert!(!self.connected_peers.contains(disconnected.preimage()));
self.queued_events.push_back(NetworkBehaviourAction::DialPeer {
peer_id: disconnected.into_preimage(),
})
},
}
} else {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::UnroutablePeer { peer }
));
}
}
},
_ => {}
}
}
fn query_finished(&mut self, q: Query<QueryInner>, params: &mut impl PollParameters)
-> Option<KademliaEvent>
{
log::trace!("Query {:?} finished.", q.id());
let result = q.into_result();
match result.inner.info {
QueryInfo::Bootstrap { peer } => {
let local_key = self.kbuckets.local_key().clone();
if &peer == local_key.preimage() {
let targets = self.kbuckets.buckets()
.skip_while(|b| b.num_entries() == 0)
.skip(1)
.map(|b| {
let mut target = kbucket::Key::new(PeerId::random());
for _ in 0 .. 16 {
let d = local_key.distance(&target);
if b.contains(&d) {
break;
}
target = kbucket::Key::new(PeerId::random());
}
target
}).collect::<Vec<_>>();
for target in targets {
let info = QueryInfo::Bootstrap { peer: target.clone().into_preimage() };
let peers = self.kbuckets.closest_keys(&target);
let inner = QueryInner::new(info);
self.queries.add_iter_closest(target.clone(), peers, inner);
}
}
Some(KademliaEvent::BootstrapResult(Ok(BootstrapOk { peer })))
}
QueryInfo::GetClosestPeers { key, .. } => {
Some(KademliaEvent::GetClosestPeersResult(Ok(
GetClosestPeersOk { key, peers: result.peers.collect() }
)))
}
QueryInfo::GetProviders { key, providers } => {
Some(KademliaEvent::GetProvidersResult(Ok(
GetProvidersOk {
key,
providers,
closest_peers: result.peers.collect()
}
)))
}
QueryInfo::PrepareAddProvider { key, context } => {
let closest_peers = result.peers.map(kbucket::Key::from);
let provider_id = params.local_peer_id().clone();
let external_addresses = params.external_addresses().collect();
let inner = QueryInner::new(QueryInfo::AddProvider {
key,
provider_id,
external_addresses,
context,
});
self.queries.add_fixed(closest_peers, inner);
None
}
QueryInfo::AddProvider { key, context, .. } => {
match context {
AddProviderContext::Publish => {
Some(KademliaEvent::StartProvidingResult(Ok(
AddProviderOk { key }
)))
}
AddProviderContext::Republish => {
Some(KademliaEvent::RepublishProviderResult(Ok(
AddProviderOk { key }
)))
}
}
}
QueryInfo::GetRecord { key, records, quorum, cache_at } => {
let result = if records.len() >= quorum.get() {
if let Some(cache_key) = cache_at {
let record = records.first().expect("[not empty]").clone();
let quorum = NonZeroUsize::new(1).expect("1 > 0");
let context = PutRecordContext::Cache;
let info = QueryInfo::PutRecord { record, quorum, context, num_results: 0 };
let inner = QueryInner::new(info);
self.queries.add_fixed(iter::once(cache_key), inner);
}
Ok(GetRecordOk { records })
} else if records.is_empty() {
Err(GetRecordError::NotFound {
key,
closest_peers: result.peers.collect()
})
} else {
Err(GetRecordError::QuorumFailed { key, records, quorum })
};
Some(KademliaEvent::GetRecordResult(result))
}
QueryInfo::PreparePutRecord { record, quorum, context } => {
let closest_peers = result.peers.map(kbucket::Key::from);
let info = QueryInfo::PutRecord { record, quorum, context, num_results: 0 };
let inner = QueryInner::new(info);
self.queries.add_fixed(closest_peers, inner);
None
}
QueryInfo::PutRecord { record, quorum, num_results, context } => {
let result = |key: record::Key| {
if num_results >= quorum.get() {
Ok(PutRecordOk { key })
} else {
Err(PutRecordError::QuorumFailed { key, quorum, num_results })
}
};
match context {
PutRecordContext::Publish =>
Some(KademliaEvent::PutRecordResult(result(record.key))),
PutRecordContext::Republish =>
Some(KademliaEvent::RepublishRecordResult(result(record.key))),
PutRecordContext::Replicate => {
debug!("Record replicated: {:?}", record.key);
None
}
PutRecordContext::Cache => {
debug!("Record cached: {:?}", record.key);
None
}
}
}
}
}
fn query_timeout(&self, query: Query<QueryInner>) -> Option<KademliaEvent> {
log::trace!("Query {:?} timed out.", query.id());
let result = query.into_result();
match result.inner.info {
QueryInfo::Bootstrap { peer } =>
Some(KademliaEvent::BootstrapResult(Err(
BootstrapError::Timeout { peer }))),
QueryInfo::PrepareAddProvider { key, context } =>
Some(match context {
AddProviderContext::Publish =>
KademliaEvent::StartProvidingResult(Err(
AddProviderError::Timeout { key })),
AddProviderContext::Republish =>
KademliaEvent::RepublishProviderResult(Err(
AddProviderError::Timeout { key })),
}),
QueryInfo::AddProvider { key, context, .. } =>
Some(match context {
AddProviderContext::Publish =>
KademliaEvent::StartProvidingResult(Err(
AddProviderError::Timeout { key })),
AddProviderContext::Republish =>
KademliaEvent::RepublishProviderResult(Err(
AddProviderError::Timeout { key })),
}),
QueryInfo::GetClosestPeers { key } => {
Some(KademliaEvent::GetClosestPeersResult(Err(
GetClosestPeersError::Timeout {
key,
peers: result.peers.collect()
})))
},
QueryInfo::PreparePutRecord { record, quorum, context, .. } => {
let err = Err(PutRecordError::Timeout {
key: record.key,
num_results: 0,
quorum
});
match context {
PutRecordContext::Publish =>
Some(KademliaEvent::PutRecordResult(err)),
PutRecordContext::Republish =>
Some(KademliaEvent::RepublishRecordResult(err)),
PutRecordContext::Replicate => {
warn!("Locating closest peers for replication failed: {:?}", err);
None
}
PutRecordContext::Cache =>
unreachable!()
}
}
QueryInfo::PutRecord { record, quorum, num_results, context } => {
let err = Err(PutRecordError::Timeout {
key: record.key,
num_results,
quorum
});
match context {
PutRecordContext::Publish =>
Some(KademliaEvent::PutRecordResult(err)),
PutRecordContext::Republish =>
Some(KademliaEvent::RepublishRecordResult(err)),
PutRecordContext::Replicate => {
debug!("Replicatiing record failed: {:?}", err);
None
}
PutRecordContext::Cache => {
debug!("Caching record failed: {:?}", err);
None
}
}
}
QueryInfo::GetRecord { key, records, quorum, .. } =>
Some(KademliaEvent::GetRecordResult(Err(
GetRecordError::Timeout { key, records, quorum }))),
QueryInfo::GetProviders { key, providers } =>
Some(KademliaEvent::GetProvidersResult(Err(
GetProvidersError::Timeout {
key,
providers,
closest_peers: result.peers.collect()
}))),
}
}
fn record_received(&mut self, source: PeerId, request_id: KademliaRequestId, mut record: Record) {
if record.publisher.as_ref() == Some(self.kbuckets.local_key().preimage()) {
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id: source,
event: KademliaHandlerIn::PutRecordRes {
key: record.key,
value: record.value,
request_id,
},
});
return
}
let now = Instant::now();
let target = kbucket::Key::new(record.key.clone());
let num_between = self.kbuckets.count_nodes_between(&target);
let k = self.queries.config().replication_factor.get();
let num_beyond_k = (usize::max(k, num_between) - k) as u32;
let expiration = self.record_ttl.map(|ttl|
now + Duration::from_secs(ttl.as_secs() >> num_beyond_k)
);
record.expires = record.expires.or(expiration).min(expiration);
if let Some(job) = self.put_record_job.as_mut() {
job.skip(record.key.clone())
}
match self.store.put(record.clone()) {
Ok(()) => {
debug!("Record stored: {:?}; {} bytes", record.key, record.value.len());
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id: source,
event: KademliaHandlerIn::PutRecordRes {
key: record.key,
value: record.value,
request_id,
},
})
}
Err(e) => {
info!("Record not stored: {:?}", e);
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id: source,
event: KademliaHandlerIn::Reset(request_id)
})
}
}
}
fn provider_received(&mut self, key: record::Key, provider: KadPeer) {
self.queued_events.push_back(NetworkBehaviourAction::GenerateEvent(
KademliaEvent::Discovered {
peer_id: provider.node_id.clone(),
addresses: provider.multiaddrs.clone(),
ty: provider.connection_ty,
}));
if &provider.node_id != self.kbuckets.local_key().preimage() {
let record = ProviderRecord {
key,
provider: provider.node_id,
expires: self.provider_record_ttl.map(|ttl| Instant::now() + ttl)
};
if let Err(e) = self.store.add_provider(record) {
info!("Provider record not stored: {:?}", e);
}
}
}
}
impl<TStore> NetworkBehaviour for Kademlia<TStore>
where
for<'a> TStore: RecordStore<'a>,
TStore: Send + 'static,
{
type ProtocolsHandler = KademliaHandler<QueryId>;
type OutEvent = KademliaEvent;
fn new_handler(&mut self) -> Self::ProtocolsHandler {
let mut handler = KademliaHandler::dial_and_listen();
if let Some(name) = self.protocol_name_override.as_ref() {
handler = handler.with_protocol_name(name.clone());
}
handler
}
fn addresses_of_peer(&mut self, peer_id: &PeerId) -> Vec<Multiaddr> {
let key = kbucket::Key::new(peer_id.clone());
let mut peer_addrs =
if let kbucket::Entry::Present(mut entry, _) = self.kbuckets.entry(&key) {
let addrs = entry.value().iter().cloned().collect::<Vec<_>>();
debug_assert!(!addrs.is_empty(), "Empty peer addresses in routing table.");
addrs
} else {
Vec::new()
};
for query in self.queries.iter() {
if let Some(addrs) = query.inner.addresses.get(peer_id) {
peer_addrs.extend(addrs.iter().cloned())
}
}
peer_addrs
}
fn inject_connected(&mut self, peer: PeerId, endpoint: ConnectedPoint) {
for (peer_id, event) in self.queries.iter_mut().filter_map(|q|
q.inner.pending_rpcs.iter()
.position(|(p, _)| p == &peer)
.map(|p| q.inner.pending_rpcs.remove(p)))
{
self.queued_events.push_back(NetworkBehaviourAction::SendEvent { peer_id, event });
}
let address = match endpoint {
ConnectedPoint::Dialer { address } => Some(address),
ConnectedPoint::Listener { .. } => None,
};
self.connection_updated(peer.clone(), address, NodeStatus::Connected);
self.connected_peers.insert(peer);
}
fn inject_addr_reach_failure(
&mut self,
peer_id: Option<&PeerId>,
addr: &Multiaddr,
err: &dyn error::Error
) {
if let Some(peer_id) = peer_id {
let key = kbucket::Key::new(peer_id.clone());
if let Some(addrs) = self.kbuckets.entry(&key).value() {
if addrs.remove(addr).is_ok() {
debug!("Address '{}' removed from peer '{}' due to error: {}.",
addr, peer_id, err);
} else {
debug!("Last remaining address '{}' of peer '{}' is unreachable: {}.",
addr, peer_id, err)
}
}
for query in self.queries.iter_mut() {
if let Some(addrs) = query.inner.addresses.get_mut(peer_id) {
addrs.retain(|a| a != addr);
}
}
}
}
fn inject_dial_failure(&mut self, peer_id: &PeerId) {
for query in self.queries.iter_mut() {
query.on_failure(peer_id);
}
}
fn inject_disconnected(&mut self, id: &PeerId, _old_endpoint: ConnectedPoint) {
for query in self.queries.iter_mut() {
query.on_failure(id);
}
self.connection_updated(id.clone(), None, NodeStatus::Disconnected);
self.connected_peers.remove(id);
}
fn inject_replaced(&mut self, peer_id: PeerId, _old: ConnectedPoint, new_endpoint: ConnectedPoint) {
for query in self.queries.iter() {
if query.is_waiting(&peer_id) {
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id: peer_id.clone(),
event: query.inner.info.to_request(query.id()),
});
}
}
if let Some(addrs) = self.kbuckets.entry(&kbucket::Key::new(peer_id)).value() {
if let ConnectedPoint::Dialer { address } = new_endpoint {
addrs.insert(address);
}
}
}
fn inject_node_event(&mut self, source: PeerId, event: KademliaHandlerEvent<QueryId>) {
match event {
KademliaHandlerEvent::FindNodeReq { key, request_id } => {
let closer_peers = self.find_closest(&kbucket::Key::new(key), &source);
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id: source,
event: KademliaHandlerIn::FindNodeRes {
closer_peers,
request_id,
},
});
}
KademliaHandlerEvent::FindNodeRes {
closer_peers,
user_data,
} => {
self.discovered(&user_data, &source, closer_peers.iter());
}
KademliaHandlerEvent::GetProvidersReq { key, request_id } => {
let provider_peers = self.provider_peers(&key, &source);
let closer_peers = self.find_closest(&kbucket::Key::new(key), &source);
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id: source,
event: KademliaHandlerIn::GetProvidersRes {
closer_peers,
provider_peers,
request_id,
},
});
}
KademliaHandlerEvent::GetProvidersRes {
closer_peers,
provider_peers,
user_data,
} => {
let peers = closer_peers.iter().chain(provider_peers.iter());
self.discovered(&user_data, &source, peers);
if let Some(query) = self.queries.get_mut(&user_data) {
if let QueryInfo::GetProviders {
providers, ..
} = &mut query.inner.info {
for peer in provider_peers {
providers.push(peer.node_id);
}
}
}
}
KademliaHandlerEvent::QueryError { user_data, error } => {
log::debug!("Request to {:?} in query {:?} failed with {:?}",
source, user_data, error);
if let Some(query) = self.queries.get_mut(&user_data) {
query.on_failure(&source)
}
}
KademliaHandlerEvent::AddProvider { key, provider } => {
if provider.node_id != source {
return
}
self.provider_received(key, provider)
}
KademliaHandlerEvent::GetRecord { key, request_id } => {
let record = match self.store.get(&key) {
Some(record) => {
if record.is_expired(Instant::now()) {
self.store.remove(&key);
None
} else {
Some(record.into_owned())
}
},
None => None
};
let closer_peers =
if record.is_none() {
self.find_closest(&kbucket::Key::new(key), &source)
} else {
Vec::new()
};
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id: source,
event: KademliaHandlerIn::GetRecordRes {
record,
closer_peers,
request_id,
},
});
}
KademliaHandlerEvent::GetRecordRes {
record,
closer_peers,
user_data,
} => {
if let Some(query) = self.queries.get_mut(&user_data) {
if let QueryInfo::GetRecord {
key, records, quorum, cache_at
} = &mut query.inner.info {
if let Some(record) = record {
records.push(record);
if records.len() == quorum.get() {
query.finish()
}
} else if quorum.get() == 1 {
let source_key = kbucket::Key::from(source.clone());
if let Some(cache_key) = cache_at {
let key = kbucket::Key::new(key.clone());
if source_key.distance(&key) < cache_key.distance(&key) {
*cache_at = Some(source_key)
}
} else {
*cache_at = Some(source_key)
}
}
}
}
self.discovered(&user_data, &source, closer_peers.iter());
}
KademliaHandlerEvent::PutRecord {
record,
request_id
} => {
self.record_received(source, request_id, record);
}
KademliaHandlerEvent::PutRecordRes {
user_data, ..
} => {
if let Some(query) = self.queries.get_mut(&user_data) {
query.on_success(&source, vec![]);
if let QueryInfo::PutRecord {
num_results, quorum, ..
} = &mut query.inner.info {
*num_results += 1;
if *num_results == quorum.get() {
query.finish()
}
}
}
}
};
}
fn poll(&mut self, cx: &mut Context, parameters: &mut impl PollParameters) -> Poll<
NetworkBehaviourAction<
<Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
Self::OutEvent,
>,
> {
let now = Instant::now();
let mut jobs_query_capacity = JOBS_MAX_QUERIES.saturating_sub(self.queries.size());
if let Some(mut job) = self.add_provider_job.take() {
let num = usize::min(JOBS_MAX_NEW_QUERIES, jobs_query_capacity);
for _ in 0 .. num {
if let Poll::Ready(r) = job.poll(cx, &mut self.store, now) {
self.start_add_provider(r.key, AddProviderContext::Republish)
} else {
break
}
}
jobs_query_capacity -= num;
self.add_provider_job = Some(job);
}
if let Some(mut job) = self.put_record_job.take() {
let num = usize::min(JOBS_MAX_NEW_QUERIES, jobs_query_capacity);
for _ in 0 .. num {
if let Poll::Ready(r) = job.poll(cx, &mut self.store, now) {
let context = if r.publisher.as_ref() == Some(self.kbuckets.local_key().preimage()) {
PutRecordContext::Republish
} else {
PutRecordContext::Replicate
};
self.start_put_record(r, Quorum::All, context)
} else {
break
}
}
self.put_record_job = Some(job);
}
loop {
if let Some(event) = self.queued_events.pop_front() {
return Poll::Ready(event);
}
if let Some(entry) = self.kbuckets.take_applied_pending() {
let kbucket::Node { key, value } = entry.inserted;
let event = KademliaEvent::RoutingUpdated {
peer: key.into_preimage(),
addresses: value,
old_peer: entry.evicted.map(|n| n.key.into_preimage())
};
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event))
}
loop {
match self.queries.poll(now) {
QueryPoolState::Finished(q) => {
if let Some(event) = self.query_finished(q, parameters) {
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event))
}
}
QueryPoolState::Timeout(q) => {
if let Some(event) = self.query_timeout(q) {
return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event))
}
}
QueryPoolState::Waiting(Some((query, peer_id))) => {
let event = query.inner.info.to_request(query.id());
if let QueryInfo::AddProvider { .. } = &query.inner.info {
query.on_success(&peer_id, vec![])
}
if self.connected_peers.contains(&peer_id) {
self.queued_events.push_back(NetworkBehaviourAction::SendEvent {
peer_id, event
});
} else if &peer_id != self.kbuckets.local_key().preimage() {
query.inner.pending_rpcs.push((peer_id.clone(), event));
self.queued_events.push_back(NetworkBehaviourAction::DialPeer {
peer_id
});
}
}
QueryPoolState::Waiting(None) | QueryPoolState::Idle => break,
}
}
if self.queued_events.is_empty() {
return Poll::Pending
}
}
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Quorum {
One,
Majority,
All,
N(NonZeroUsize)
}
impl Quorum {
fn eval(&self, total: NonZeroUsize) -> NonZeroUsize {
match self {
Quorum::One => NonZeroUsize::new(1).expect("1 != 0"),
Quorum::Majority => NonZeroUsize::new(total.get() / 2 + 1).expect("n + 1 != 0"),
Quorum::All => total,
Quorum::N(n) => NonZeroUsize::min(total, *n)
}
}
}
#[derive(Debug)]
pub enum KademliaEvent {
BootstrapResult(BootstrapResult),
GetClosestPeersResult(GetClosestPeersResult),
GetProvidersResult(GetProvidersResult),
StartProvidingResult(AddProviderResult),
RepublishProviderResult(AddProviderResult),
GetRecordResult(GetRecordResult),
PutRecordResult(PutRecordResult),
RepublishRecordResult(PutRecordResult),
Discovered {
peer_id: PeerId,
addresses: Vec<Multiaddr>,
ty: KadConnectionType,
},
RoutingUpdated {
peer: PeerId,
addresses: Addresses,
old_peer: Option<PeerId>,
},
UnroutablePeer {
peer: PeerId
}
}
pub type GetRecordResult = Result<GetRecordOk, GetRecordError>;
#[derive(Debug, Clone)]
pub struct GetRecordOk {
pub records: Vec<Record>
}
#[derive(Debug, Clone)]
pub enum GetRecordError {
NotFound {
key: record::Key,
closest_peers: Vec<PeerId>
},
QuorumFailed {
key: record::Key,
records: Vec<Record>,
quorum: NonZeroUsize
},
Timeout {
key: record::Key,
records: Vec<Record>,
quorum: NonZeroUsize
}
}
impl GetRecordError {
pub fn key(&self) -> &record::Key {
match self {
GetRecordError::QuorumFailed { key, .. } => key,
GetRecordError::Timeout { key, .. } => key,
GetRecordError::NotFound { key, .. } => key,
}
}
pub fn into_key(self) -> record::Key {
match self {
GetRecordError::QuorumFailed { key, .. } => key,
GetRecordError::Timeout { key, .. } => key,
GetRecordError::NotFound { key, .. } => key,
}
}
}
pub type PutRecordResult = Result<PutRecordOk, PutRecordError>;
#[derive(Debug, Clone)]
pub struct PutRecordOk {
pub key: record::Key
}
#[derive(Debug)]
pub enum PutRecordError {
QuorumFailed {
key: record::Key,
num_results: usize,
quorum: NonZeroUsize
},
Timeout {
key: record::Key,
num_results: usize,
quorum: NonZeroUsize
},
LocalStorageError {
key: record::Key,
cause: store::Error
}
}
impl PutRecordError {
pub fn key(&self) -> &record::Key {
match self {
PutRecordError::QuorumFailed { key, .. } => key,
PutRecordError::Timeout { key, .. } => key,
PutRecordError::LocalStorageError { key, .. } => key
}
}
pub fn into_key(self) -> record::Key {
match self {
PutRecordError::QuorumFailed { key, .. } => key,
PutRecordError::Timeout { key, .. } => key,
PutRecordError::LocalStorageError { key, .. } => key,
}
}
}
pub type BootstrapResult = Result<BootstrapOk, BootstrapError>;
#[derive(Debug, Clone)]
pub struct BootstrapOk {
pub peer: PeerId
}
#[derive(Debug, Clone)]
pub enum BootstrapError {
Timeout { peer: PeerId }
}
pub type GetClosestPeersResult = Result<GetClosestPeersOk, GetClosestPeersError>;
#[derive(Debug, Clone)]
pub struct GetClosestPeersOk {
pub key: Vec<u8>,
pub peers: Vec<PeerId>
}
#[derive(Debug, Clone)]
pub enum GetClosestPeersError {
Timeout {
key: Vec<u8>,
peers: Vec<PeerId>
}
}
impl GetClosestPeersError {
pub fn key(&self) -> &Vec<u8> {
match self {
GetClosestPeersError::Timeout { key, .. } => key,
}
}
pub fn into_key(self) -> Vec<u8> {
match self {
GetClosestPeersError::Timeout { key, .. } => key,
}
}
}
pub type GetProvidersResult = Result<GetProvidersOk, GetProvidersError>;
#[derive(Debug, Clone)]
pub struct GetProvidersOk {
pub key: record::Key,
pub providers: Vec<PeerId>,
pub closest_peers: Vec<PeerId>
}
#[derive(Debug, Clone)]
pub enum GetProvidersError {
Timeout {
key: record::Key,
providers: Vec<PeerId>,
closest_peers: Vec<PeerId>
}
}
impl GetProvidersError {
pub fn key(&self) -> &record::Key {
match self {
GetProvidersError::Timeout { key, .. } => key,
}
}
pub fn into_key(self) -> record::Key {
match self {
GetProvidersError::Timeout { key, .. } => key,
}
}
}
pub type AddProviderResult = Result<AddProviderOk, AddProviderError>;
#[derive(Debug, Clone)]
pub struct AddProviderOk {
pub key: record::Key,
}
#[derive(Debug)]
pub enum AddProviderError {
Timeout {
key: record::Key,
},
LocalStorageError {
key: record::Key,
cause: store::Error
}
}
impl AddProviderError {
pub fn key(&self) -> &record::Key {
match self {
AddProviderError::Timeout { key, .. } => key,
AddProviderError::LocalStorageError { key, .. } => key,
}
}
pub fn into_key(self) -> record::Key {
match self {
AddProviderError::Timeout { key, .. } => key,
AddProviderError::LocalStorageError { key, .. } => key,
}
}
}
impl From<kbucket::EntryView<kbucket::Key<PeerId>, Addresses>> for KadPeer {
fn from(e: kbucket::EntryView<kbucket::Key<PeerId>, Addresses>) -> KadPeer {
KadPeer {
node_id: e.node.key.into_preimage(),
multiaddrs: e.node.value.into_vec(),
connection_ty: match e.status {
NodeStatus::Connected => KadConnectionType::Connected,
NodeStatus::Disconnected => KadConnectionType::NotConnected
}
}
}
}
struct QueryInner {
info: QueryInfo,
addresses: FnvHashMap<PeerId, SmallVec<[Multiaddr; 8]>>,
pending_rpcs: SmallVec<[(PeerId, KademliaHandlerIn<QueryId>); K_VALUE.get()]>
}
impl QueryInner {
fn new(info: QueryInfo) -> Self {
QueryInner {
info,
addresses: Default::default(),
pending_rpcs: SmallVec::default()
}
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum AddProviderContext {
Publish,
Republish,
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum PutRecordContext {
Publish,
Republish,
Replicate,
Cache,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum QueryInfo {
Bootstrap {
peer: PeerId,
},
GetClosestPeers { key: Vec<u8> },
GetProviders {
key: record::Key,
providers: Vec<PeerId>,
},
PrepareAddProvider {
key: record::Key,
context: AddProviderContext,
},
AddProvider {
key: record::Key,
provider_id: PeerId,
external_addresses: Vec<Multiaddr>,
context: AddProviderContext,
},
PreparePutRecord {
record: Record,
quorum: NonZeroUsize,
context: PutRecordContext,
},
PutRecord {
record: Record,
quorum: NonZeroUsize,
num_results: usize,
context: PutRecordContext,
},
GetRecord {
key: record::Key,
records: Vec<Record>,
quorum: NonZeroUsize,
cache_at: Option<kbucket::Key<PeerId>>,
},
}
impl QueryInfo {
fn to_request(&self, query_id: QueryId) -> KademliaHandlerIn<QueryId> {
match &self {
QueryInfo::Bootstrap { peer } => KademliaHandlerIn::FindNodeReq {
key: peer.clone().into_bytes(),
user_data: query_id,
},
QueryInfo::GetClosestPeers { key, .. } => KademliaHandlerIn::FindNodeReq {
key: key.clone(),
user_data: query_id,
},
QueryInfo::GetProviders { key, .. } => KademliaHandlerIn::GetProvidersReq {
key: key.clone(),
user_data: query_id,
},
QueryInfo::PrepareAddProvider { key, .. } => KademliaHandlerIn::FindNodeReq {
key: key.to_vec(),
user_data: query_id,
},
QueryInfo::AddProvider {
key,
provider_id,
external_addresses,
..
} => KademliaHandlerIn::AddProvider {
key: key.clone(),
provider: crate::protocol::KadPeer {
node_id: provider_id.clone(),
multiaddrs: external_addresses.clone(),
connection_ty: crate::protocol::KadConnectionType::Connected,
}
},
QueryInfo::GetRecord { key, .. } => KademliaHandlerIn::GetRecord {
key: key.clone(),
user_data: query_id,
},
QueryInfo::PreparePutRecord { record, .. } => KademliaHandlerIn::FindNodeReq {
key: record.key.to_vec(),
user_data: query_id,
},
QueryInfo::PutRecord { record, .. } => KademliaHandlerIn::PutRecord {
record: record.clone(),
user_data: query_id
}
}
}
}