validator_discovery: less flexible, but simpler design (#3052)

* validator_discovery: less flexible, but simpler design

* fix test

* remove unused struct

* smol optimization
This commit is contained in:
Andronik Ordian
2021-05-19 18:54:13 +02:00
committed by GitHub
parent 44c03a3633
commit 98c06f5b57
7 changed files with 84 additions and 422 deletions
-2
View File
@@ -509,7 +509,6 @@ where
NetworkBridgeMessage::ConnectToValidators {
validator_ids,
peer_set,
keep_alive,
} => {
tracing::trace!(
target: LOG_TARGET,
@@ -522,7 +521,6 @@ where
let (ns, ads) = validator_discovery.on_request(
validator_ids,
peer_set,
keep_alive,
network_service,
authority_discovery_service,
).await;
@@ -19,10 +19,9 @@
use crate::Network;
use core::marker::PhantomData;
use std::collections::{HashSet, HashMap, hash_map};
use std::collections::HashSet;
use async_trait::async_trait;
use futures::channel::oneshot;
use sc_network::multiaddr::Multiaddr;
use sc_authority_discovery::Service as AuthorityDiscoveryService;
@@ -52,51 +51,6 @@ impl AuthorityDiscovery for AuthorityDiscoveryService {
}
}
/// This struct tracks the state for one `ConnectToValidators` request.
struct NonRevokedConnectionRequestState {
requested: Vec<AuthorityDiscoveryId>,
keep_alive: oneshot::Receiver<()>,
}
impl NonRevokedConnectionRequestState {
/// Create a new instance of `ConnectToValidatorsState`.
pub fn new(
requested: Vec<AuthorityDiscoveryId>,
keep_alive: oneshot::Receiver<()>,
) -> Self {
Self {
requested,
keep_alive,
}
}
/// Returns `true` if the request is revoked.
pub fn is_revoked(&mut self) -> bool {
self.keep_alive.try_recv().is_err()
}
pub fn requested(&self) -> &[AuthorityDiscoveryId] {
self.requested.as_ref()
}
}
/// Will be called by [`Service::on_request`] when a request was revoked.
///
/// Takes the `map` of requested validators and the `id` of the validator that should be revoked.
///
/// Returns `Some(id)` iff the request counter is `0`.
fn on_revoke(map: &mut HashMap<AuthorityDiscoveryId, u64>, id: AuthorityDiscoveryId) -> Option<AuthorityDiscoveryId> {
if let hash_map::Entry::Occupied(mut entry) = map.entry(id) {
*entry.get_mut() = entry.get().saturating_sub(1);
if *entry.get() == 0 {
return Some(entry.remove_entry().0);
}
}
None
}
pub(super) struct Service<N, AD> {
state: PerPeerSet<StatePerPeerSet>,
// PhantomData used to make the struct generic instead of having generic methods
@@ -105,111 +59,67 @@ pub(super) struct Service<N, AD> {
#[derive(Default)]
struct StatePerPeerSet {
// The `u64` counts the number of pending non-revoked requests for this validator
// note: the validators in this map are not necessarily present
// in the `connected_validators` map.
// Invariant: the value > 0 for non-revoked requests.
requested_validators: HashMap<AuthorityDiscoveryId, u64>,
non_revoked_discovery_requests: Vec<NonRevokedConnectionRequestState>,
previously_requested: HashSet<Multiaddr>,
}
impl<N: Network, AD: AuthorityDiscovery> Service<N, AD> {
pub fn new() -> Self {
Self {
state: PerPeerSet::default(),
state: Default::default(),
_phantom: PhantomData,
}
}
/// On a new connection request, a peer set update will be issued.
/// It will ask the network to connect to the validators and not disconnect
/// from them at least until all the pending requests containing them are revoked.
/// from them at least until the next request is issued for the same peer set.
///
/// This method will also clean up all previously revoked requests.
/// This method will also disconnect from previously connected validators not in the `validator_ids` set.
/// it takes `network_service` and `authority_discovery_service` by value
/// and returns them as a workaround for the Future: Send requirement imposed by async fn impl.
pub async fn on_request(
&mut self,
validator_ids: Vec<AuthorityDiscoveryId>,
peer_set: PeerSet,
keep_alive: oneshot::Receiver<()>,
mut network_service: N,
mut authority_discovery_service: AD,
) -> (N, AD) {
const MAX_ADDR_PER_PEER: usize = 3;
let state = &mut self.state[peer_set];
// Increment the counter of how many times the validators were requested.
validator_ids.iter().for_each(|id| *state.requested_validators.entry(id.clone()).or_default() += 1);
// collect multiaddress of validators
let mut multiaddr_to_add = HashSet::new();
for authority in validator_ids.iter() {
let mut newly_requested = HashSet::new();
for authority in validator_ids.into_iter() {
let result = authority_discovery_service.get_addresses_by_authority_id(authority.clone()).await;
if let Some(addresses) = result {
// We might have several `PeerId`s per `AuthorityId`
multiaddr_to_add.extend(addresses.into_iter().take(MAX_ADDR_PER_PEER));
newly_requested.extend(addresses);
} else {
tracing::debug!(target: LOG_TARGET, "Authority Discovery couldn't resolve {:?}", authority);
}
}
let state = &mut self.state[peer_set];
// clean up revoked requests
let mut revoked_indices = Vec::new();
let mut revoked_validators = Vec::new();
for (i, maybe_revoked) in state.non_revoked_discovery_requests.iter_mut().enumerate() {
if maybe_revoked.is_revoked() {
for id in maybe_revoked.requested() {
if let Some(id) = on_revoke(&mut state.requested_validators, id.clone()) {
revoked_validators.push(id);
}
}
revoked_indices.push(i);
}
}
// clean up revoked requests states
//
// note that the `.rev()` here is important to guarantee `swap_remove`
// doesn't invalidate unprocessed `revoked_indices`
for to_revoke in revoked_indices.into_iter().rev() {
drop(state.non_revoked_discovery_requests.swap_remove(to_revoke));
}
// multiaddresses to remove
let mut multiaddr_to_remove = HashSet::new();
for id in revoked_validators.into_iter() {
let result = authority_discovery_service.get_addresses_by_authority_id(id.clone()).await;
if let Some(addresses) = result {
multiaddr_to_remove.extend(addresses.into_iter());
} else {
tracing::debug!(
target: LOG_TARGET,
"Authority Discovery couldn't resolve {:?} on cleanup, a leak is possible",
id,
);
}
}
let multiaddr_to_remove = state.previously_requested
.difference(&newly_requested)
.cloned()
.collect();
let multiaddr_to_add = newly_requested.difference(&state.previously_requested)
.cloned()
.collect();
state.previously_requested = newly_requested;
// ask the network to connect to these nodes and not disconnect
// from them until removed from the set
if let Err(e) = network_service.add_to_peers_set(
peer_set.into_protocol_name(),
multiaddr_to_add.clone(),
multiaddr_to_add,
).await {
tracing::warn!(target: LOG_TARGET, err = ?e, "AuthorityDiscoveryService returned an invalid multiaddress");
}
// the addresses are known to be valid
let _ = network_service.remove_from_peers_set(
peer_set.into_protocol_name(),
multiaddr_to_remove.clone()
multiaddr_to_remove
).await;
state.non_revoked_discovery_requests.push(NonRevokedConnectionRequestState::new(
validator_ids,
keep_alive,
));
(network_service, authority_discovery_service)
}
}
@@ -219,7 +129,7 @@ mod tests {
use super::*;
use crate::network::{Network, NetworkAction};
use std::{borrow::Cow, pin::Pin};
use std::{borrow::Cow, pin::Pin, collections::HashMap};
use futures::{sink::Sink, stream::BoxStream};
use sc_network::{Event as NetworkEvent, IfDisconnected};
use sp_keyring::Sr25519Keyring;
@@ -317,26 +227,9 @@ mod tests {
"/ip4/127.0.0.1/tcp/1236".parse().unwrap(),
]
}
#[test]
fn request_is_revoked_when_the_receiver_is_dropped() {
let (keep_alive_handle, keep_alive) = oneshot::channel();
let mut request = NonRevokedConnectionRequestState::new(
Vec::new(),
keep_alive,
);
assert!(!request.is_revoked());
drop(keep_alive_handle);
assert!(request.is_revoked());
}
// Test cleanup works.
#[test]
fn requests_are_removed_on_revoke() {
fn old_multiaddrs_are_removed_on_new_request() {
let mut service = new_service();
let (ns, ads) = new_network();
@@ -344,87 +237,22 @@ mod tests {
let authority_ids: Vec<_> = ads.by_peer_id.values().cloned().collect();
futures::executor::block_on(async move {
let (keep_alive_handle, keep_alive) = oneshot::channel();
let (ns, ads) = service.on_request(
vec![authority_ids[0].clone()],
PeerSet::Validation,
keep_alive,
ns,
ads,
).await;
// revoke the request
drop(keep_alive_handle);
let (_keep_alive_handle, keep_alive) = oneshot::channel();
let _ = service.on_request(
vec![authority_ids[1].clone()],
PeerSet::Validation,
keep_alive,
ns,
ads,
).await;
let state = &service.state[PeerSet::Validation];
assert_eq!(state.non_revoked_discovery_requests.len(), 1);
});
}
// More complex test with overlapping revoked requests
#[test]
fn revoking_requests_with_overlapping_validator_sets() {
let mut service = new_service();
let (ns, ads) = new_network();
let authority_ids: Vec<_> = ads.by_peer_id.values().cloned().collect();
futures::executor::block_on(async move {
let (keep_alive_handle, keep_alive) = oneshot::channel();
let (ns, ads) = service.on_request(
vec![authority_ids[0].clone(), authority_ids[2].clone()],
PeerSet::Validation,
keep_alive,
ns,
ads,
).await;
// revoke the first request
drop(keep_alive_handle);
let (keep_alive_handle, keep_alive) = oneshot::channel();
let (ns, ads) = service.on_request(
vec![authority_ids[0].clone(), authority_ids[1].clone()],
PeerSet::Validation,
keep_alive,
ns,
ads,
).await;
let state = &service.state[PeerSet::Validation];
assert_eq!(state.non_revoked_discovery_requests.len(), 1);
assert_eq!(ns.peers_set.len(), 2);
// revoke the second request
drop(keep_alive_handle);
let (_keep_alive_handle, keep_alive) = oneshot::channel();
let (ns, _) = service.on_request(
vec![authority_ids[0].clone()],
PeerSet::Validation,
keep_alive,
ns,
ads,
).await;
let state = &service.state[PeerSet::Validation];
assert_eq!(state.non_revoked_discovery_requests.len(), 1);
assert_eq!(ns.peers_set.len(), 1);
assert_eq!(state.previously_requested.len(), 1);
});
}
}