use crate::{
state_machine::{ConsensusGossip, TopicNotification, PERIODIC_MAINTENANCE_INTERVAL},
Network, Syncing, Validator,
};
use sc_network::{
service::traits::{NotificationEvent, ValidationResult},
types::ProtocolName,
NotificationService, ReputationChange,
};
use sc_network_sync::SyncEvent;
use futures::{
channel::mpsc::{channel, Receiver, Sender},
prelude::*,
};
use log::trace;
use prometheus_endpoint::Registry;
use sc_network_types::PeerId;
use sp_runtime::traits::Block as BlockT;
use std::{
collections::{HashMap, VecDeque},
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
pub struct GossipEngine<B: BlockT> {
state_machine: ConsensusGossip<B>,
network: Box<dyn Network<B> + Send>,
sync: Box<dyn Syncing<B>>,
periodic_maintenance_interval: futures_timer::Delay,
protocol: ProtocolName,
sync_event_stream: Pin<Box<dyn Stream<Item = SyncEvent> + Send>>,
notification_service: Box<dyn NotificationService>,
message_sinks: HashMap<B::Hash, Vec<Sender<TopicNotification>>>,
forwarding_state: ForwardingState<B>,
is_terminated: bool,
}
enum ForwardingState<B: BlockT> {
Idle,
Busy(VecDeque<(B::Hash, TopicNotification)>),
}
impl<B: BlockT> Unpin for GossipEngine<B> {}
impl<B: BlockT> GossipEngine<B> {
pub fn new<N, S>(
network: N,
sync: S,
notification_service: Box<dyn NotificationService>,
protocol: impl Into<ProtocolName>,
validator: Arc<dyn Validator<B>>,
metrics_registry: Option<&Registry>,
) -> Self
where
B: 'static,
N: Network<B> + Send + Clone + 'static,
S: Syncing<B> + Send + Clone + 'static,
{
let protocol = protocol.into();
let sync_event_stream = sync.event_stream("network-gossip");
GossipEngine {
state_machine: ConsensusGossip::new(validator, protocol.clone(), metrics_registry),
network: Box::new(network),
sync: Box::new(sync),
notification_service,
periodic_maintenance_interval: futures_timer::Delay::new(PERIODIC_MAINTENANCE_INTERVAL),
protocol,
sync_event_stream,
message_sinks: HashMap::new(),
forwarding_state: ForwardingState::Idle,
is_terminated: false,
}
}
pub fn report(&self, who: PeerId, reputation: ReputationChange) {
self.network.report_peer(who, reputation);
}
pub fn register_gossip_message(&mut self, topic: B::Hash, message: Vec<u8>) {
self.state_machine.register_message(topic, message);
}
pub fn broadcast_topic(&mut self, topic: B::Hash, force: bool) {
self.state_machine.broadcast_topic(&mut self.notification_service, topic, force);
}
pub fn messages_for(&mut self, topic: B::Hash) -> Receiver<TopicNotification> {
let past_messages = self.state_machine.messages_for(topic).collect::<Vec<_>>();
let (mut tx, rx) = channel(usize::max(past_messages.len(), 10));
for notification in past_messages {
tx.try_send(notification)
.expect("receiver known to be live, and buffer size known to suffice; qed");
}
self.message_sinks.entry(topic).or_default().push(tx);
rx
}
pub fn send_topic(&mut self, who: &PeerId, topic: B::Hash, force: bool) {
self.state_machine.send_topic(&mut self.notification_service, who, topic, force)
}
pub fn gossip_message(&mut self, topic: B::Hash, message: Vec<u8>, force: bool) {
self.state_machine
.multicast(&mut self.notification_service, topic, message, force)
}
pub fn send_message(&mut self, who: Vec<PeerId>, data: Vec<u8>) {
for who in &who {
self.state_machine
.send_message(&mut self.notification_service, who, data.clone());
}
}
pub fn announce(&self, block: B::Hash, associated_data: Option<Vec<u8>>) {
self.sync.announce_block(block, associated_data);
}
pub fn take_notification_service(self) -> Box<dyn NotificationService> {
self.notification_service
}
}
impl<B: BlockT> Future for GossipEngine<B> {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let this = &mut *self;
'outer: loop {
match &mut this.forwarding_state {
ForwardingState::Idle => {
let next_notification_event =
this.notification_service.next_event().poll_unpin(cx);
let sync_event_stream = this.sync_event_stream.poll_next_unpin(cx);
if next_notification_event.is_pending() && sync_event_stream.is_pending() {
break
}
match next_notification_event {
Poll::Ready(Some(event)) => match event {
NotificationEvent::ValidateInboundSubstream {
peer,
handshake,
result_tx,
..
} => {
let result = this
.network
.peer_role(peer, handshake)
.map_or(ValidationResult::Reject, |_| ValidationResult::Accept);
let _ = result_tx.send(result);
},
NotificationEvent::NotificationStreamOpened {
peer, handshake, ..
} => {
let Some(role) = this.network.peer_role(peer, handshake) else {
log::debug!(target: "gossip", "role for {peer} couldn't be determined");
continue
};
this.state_machine.new_peer(
&mut this.notification_service,
peer,
role,
);
},
NotificationEvent::NotificationStreamClosed { peer } => {
this.state_machine
.peer_disconnected(&mut this.notification_service, peer);
},
NotificationEvent::NotificationReceived { peer, notification } => {
let to_forward = this.state_machine.on_incoming(
&mut *this.network,
&mut this.notification_service,
peer,
vec![notification],
);
this.forwarding_state = ForwardingState::Busy(to_forward.into());
},
},
Poll::Ready(None) => {
self.is_terminated = true;
return Poll::Ready(())
},
Poll::Pending => {},
}
match sync_event_stream {
Poll::Ready(Some(event)) => match event {
SyncEvent::PeerConnected(remote) =>
this.network.add_set_reserved(remote, this.protocol.clone()),
SyncEvent::PeerDisconnected(remote) =>
this.network.remove_set_reserved(remote, this.protocol.clone()),
},
Poll::Ready(None) => {
self.is_terminated = true;
return Poll::Ready(())
},
Poll::Pending => {},
}
},
ForwardingState::Busy(to_forward) => {
let (topic, notification) = match to_forward.pop_front() {
Some(n) => n,
None => {
this.forwarding_state = ForwardingState::Idle;
continue
},
};
let sinks = match this.message_sinks.get_mut(&topic) {
Some(sinks) => sinks,
None => continue,
};
for sink in sinks.iter_mut() {
match sink.poll_ready(cx) {
Poll::Ready(Ok(())) => {},
Poll::Ready(Err(_)) => {},
Poll::Pending => {
to_forward.push_front((topic, notification));
break 'outer
},
}
}
sinks.retain(|sink| !sink.is_closed()); if sinks.is_empty() {
this.message_sinks.remove(&topic);
continue
}
trace!(
target: "gossip",
"Pushing consensus message to sinks for {}.", topic,
);
for sink in sinks {
match sink.start_send(notification.clone()) {
Ok(()) => {},
Err(e) if e.is_full() => {
unreachable!("Previously ensured that all sinks are ready; qed.")
},
Err(_) => {},
}
}
},
}
}
while let Poll::Ready(()) = this.periodic_maintenance_interval.poll_unpin(cx) {
this.periodic_maintenance_interval.reset(PERIODIC_MAINTENANCE_INTERVAL);
this.state_machine.tick(&mut this.notification_service);
this.message_sinks.retain(|_, sinks| {
sinks.retain(|sink| !sink.is_closed());
!sinks.is_empty()
});
}
Poll::Pending
}
}
impl<B: BlockT> futures::future::FusedFuture for GossipEngine<B> {
fn is_terminated(&self) -> bool {
self.is_terminated
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{ValidationResult, ValidatorContext};
use codec::{DecodeAll, Encode};
use futures::{
channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender},
executor::{block_on, block_on_stream},
future::poll_fn,
};
use quickcheck::{Arbitrary, Gen, QuickCheck};
use sc_network::{
config::MultiaddrWithPeerId,
service::traits::{Direction, MessageSink, NotificationEvent},
Event, NetworkBlock, NetworkEventStream, NetworkPeers, NotificationService, Roles,
};
use sc_network_common::role::ObservedRole;
use sc_network_sync::SyncEventStream;
use sc_network_types::multiaddr::Multiaddr;
use sp_runtime::{
testing::H256,
traits::{Block as BlockT, NumberFor},
};
use std::{
collections::HashSet,
sync::{Arc, Mutex},
};
use substrate_test_runtime_client::runtime::Block;
#[derive(Clone, Default)]
struct TestNetwork {}
#[async_trait::async_trait]
impl NetworkPeers for TestNetwork {
fn set_authorized_peers(&self, _peers: HashSet<PeerId>) {
unimplemented!();
}
fn set_authorized_only(&self, _reserved_only: bool) {
unimplemented!();
}
fn add_known_address(&self, _peer_id: PeerId, _addr: Multiaddr) {
unimplemented!();
}
fn report_peer(&self, _peer_id: PeerId, _cost_benefit: ReputationChange) {}
fn peer_reputation(&self, _peer_id: &PeerId) -> i32 {
unimplemented!()
}
fn disconnect_peer(&self, _peer_id: PeerId, _protocol: ProtocolName) {
unimplemented!();
}
fn accept_unreserved_peers(&self) {
unimplemented!();
}
fn deny_unreserved_peers(&self) {
unimplemented!();
}
fn add_reserved_peer(&self, _peer: MultiaddrWithPeerId) -> Result<(), String> {
unimplemented!();
}
fn remove_reserved_peer(&self, _peer_id: PeerId) {
unimplemented!();
}
fn set_reserved_peers(
&self,
_protocol: ProtocolName,
_peers: HashSet<Multiaddr>,
) -> Result<(), String> {
unimplemented!();
}
fn add_peers_to_reserved_set(
&self,
_protocol: ProtocolName,
_peers: HashSet<Multiaddr>,
) -> Result<(), String> {
unimplemented!();
}
fn remove_peers_from_reserved_set(
&self,
_protocol: ProtocolName,
_peers: Vec<PeerId>,
) -> Result<(), String> {
unimplemented!();
}
fn sync_num_connected(&self) -> usize {
unimplemented!();
}
fn peer_role(&self, _peer_id: PeerId, handshake: Vec<u8>) -> Option<ObservedRole> {
Roles::decode_all(&mut &handshake[..])
.ok()
.and_then(|role| Some(ObservedRole::from(role)))
}
async fn reserved_peers(&self) -> Result<Vec<PeerId>, ()> {
unimplemented!();
}
}
impl NetworkEventStream for TestNetwork {
fn event_stream(&self, _name: &'static str) -> Pin<Box<dyn Stream<Item = Event> + Send>> {
unimplemented!();
}
}
impl NetworkBlock<<Block as BlockT>::Hash, NumberFor<Block>> for TestNetwork {
fn announce_block(&self, _hash: <Block as BlockT>::Hash, _data: Option<Vec<u8>>) {
unimplemented!();
}
fn new_best_block_imported(
&self,
_hash: <Block as BlockT>::Hash,
_number: NumberFor<Block>,
) {
unimplemented!();
}
}
#[derive(Clone, Default)]
struct TestSync {
inner: Arc<Mutex<TestSyncInner>>,
}
#[derive(Clone, Default)]
struct TestSyncInner {
event_senders: Vec<UnboundedSender<SyncEvent>>,
}
impl SyncEventStream for TestSync {
fn event_stream(
&self,
_name: &'static str,
) -> Pin<Box<dyn Stream<Item = SyncEvent> + Send>> {
let (tx, rx) = unbounded();
self.inner.lock().unwrap().event_senders.push(tx);
Box::pin(rx)
}
}
impl NetworkBlock<<Block as BlockT>::Hash, NumberFor<Block>> for TestSync {
fn announce_block(&self, _hash: <Block as BlockT>::Hash, _data: Option<Vec<u8>>) {
unimplemented!();
}
fn new_best_block_imported(
&self,
_hash: <Block as BlockT>::Hash,
_number: NumberFor<Block>,
) {
unimplemented!();
}
}
#[derive(Debug)]
pub(crate) struct TestNotificationService {
rx: UnboundedReceiver<NotificationEvent>,
}
#[async_trait::async_trait]
impl sc_network::service::traits::NotificationService for TestNotificationService {
async fn open_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
unimplemented!();
}
async fn close_substream(&mut self, _peer: PeerId) -> Result<(), ()> {
unimplemented!();
}
fn send_sync_notification(&mut self, _peer: &PeerId, _notification: Vec<u8>) {
unimplemented!();
}
async fn send_async_notification(
&mut self,
_peer: &PeerId,
_notification: Vec<u8>,
) -> Result<(), sc_network::error::Error> {
unimplemented!();
}
async fn set_handshake(&mut self, _handshake: Vec<u8>) -> Result<(), ()> {
unimplemented!();
}
fn try_set_handshake(&mut self, _handshake: Vec<u8>) -> Result<(), ()> {
unimplemented!();
}
async fn next_event(&mut self) -> Option<NotificationEvent> {
self.rx.next().await
}
fn clone(&mut self) -> Result<Box<dyn NotificationService>, ()> {
unimplemented!();
}
fn protocol(&self) -> &ProtocolName {
unimplemented!();
}
fn message_sink(&self, _peer: &PeerId) -> Option<Box<dyn MessageSink>> {
unimplemented!();
}
}
struct AllowAll;
impl Validator<Block> for AllowAll {
fn validate(
&self,
_context: &mut dyn ValidatorContext<Block>,
_sender: &PeerId,
_data: &[u8],
) -> ValidationResult<H256> {
ValidationResult::ProcessAndKeep(H256::default())
}
}
#[test]
fn returns_when_network_event_stream_closes() {
let network = TestNetwork::default();
let sync = Arc::new(TestSync::default());
let (tx, rx) = unbounded();
let notification_service = Box::new(TestNotificationService { rx });
let mut gossip_engine = GossipEngine::<Block>::new(
network.clone(),
sync,
notification_service,
"/my_protocol",
Arc::new(AllowAll {}),
None,
);
drop(tx);
block_on(poll_fn(move |ctx| {
if let Poll::Pending = gossip_engine.poll_unpin(ctx) {
panic!(
"Expected gossip engine to finish on first poll, given that \
`GossipEngine.network_event_stream` closes right away."
)
}
Poll::Ready(())
}))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn keeps_multiple_subscribers_per_topic_updated_with_both_old_and_new_messages() {
let topic = H256::default();
let protocol = ProtocolName::from("/my_protocol");
let remote_peer = PeerId::random();
let network = TestNetwork::default();
let sync = Arc::new(TestSync::default());
let (mut tx, rx) = unbounded();
let notification_service = Box::new(TestNotificationService { rx });
let mut gossip_engine = GossipEngine::<Block>::new(
network.clone(),
sync.clone(),
notification_service,
protocol.clone(),
Arc::new(AllowAll {}),
None,
);
tx.send(NotificationEvent::NotificationStreamOpened {
peer: remote_peer,
direction: Direction::Inbound,
negotiated_fallback: None,
handshake: Roles::FULL.encode(),
})
.await
.unwrap();
let messages = vec![vec![1], vec![2]];
tx.send(NotificationEvent::NotificationReceived {
peer: remote_peer,
notification: messages[0].clone().into(),
})
.await
.unwrap();
let mut subscribers = vec![];
for _ in 0..2 {
subscribers.push(gossip_engine.messages_for(topic));
}
tx.send(NotificationEvent::NotificationReceived {
peer: remote_peer,
notification: messages[1].clone().into(),
})
.await
.unwrap();
tokio::spawn(gossip_engine);
let mut subscribers =
subscribers.into_iter().map(|s| block_on_stream(s)).collect::<Vec<_>>();
for message in messages {
for subscriber in subscribers.iter_mut() {
assert_eq!(
subscriber.next(),
Some(TopicNotification { message: message.clone(), sender: Some(remote_peer) }),
);
}
}
}
#[test]
fn forwarding_to_different_size_and_topic_channels() {
#[derive(Clone, Debug)]
struct ChannelLengthAndTopic {
length: usize,
topic: H256,
}
impl Arbitrary for ChannelLengthAndTopic {
fn arbitrary(g: &mut Gen) -> Self {
let possible_length = (0..100).collect::<Vec<usize>>();
let possible_topics = (0..10).collect::<Vec<u64>>();
Self {
length: *g.choose(&possible_length).unwrap(),
topic: H256::from_low_u64_ne(*g.choose(&possible_topics).unwrap()),
}
}
}
#[derive(Clone, Debug)]
struct Message {
topic: H256,
}
impl Arbitrary for Message {
fn arbitrary(g: &mut Gen) -> Self {
let possible_topics = (0..10).collect::<Vec<u64>>();
Self {
topic: H256::from_low_u64_ne(*g.choose(&possible_topics).unwrap()),
}
}
}
struct TestValidator;
impl Validator<Block> for TestValidator {
fn validate(
&self,
_context: &mut dyn ValidatorContext<Block>,
_sender: &PeerId,
data: &[u8],
) -> ValidationResult<H256> {
ValidationResult::ProcessAndKeep(H256::from_slice(&data[0..32]))
}
}
fn prop(channels: Vec<ChannelLengthAndTopic>, notifications: Vec<Vec<Message>>) {
let protocol = ProtocolName::from("/my_protocol");
let remote_peer = PeerId::random();
let network = TestNetwork::default();
let sync = Arc::new(TestSync::default());
let (mut tx, rx) = unbounded();
let notification_service = Box::new(TestNotificationService { rx });
let num_channels_per_topic = channels.iter().fold(
HashMap::new(),
|mut acc, ChannelLengthAndTopic { topic, .. }| {
acc.entry(topic).and_modify(|e| *e += 1).or_insert(1);
acc
},
);
let expected_msgs_per_topic_all_chan = notifications
.iter()
.fold(HashMap::new(), |mut acc, messages| {
for message in messages {
acc.entry(message.topic).and_modify(|e| *e += 1).or_insert(1);
}
acc
})
.into_iter()
.map(|(topic, num)| (topic, num_channels_per_topic.get(&topic).unwrap_or(&0) * num))
.collect::<HashMap<H256, _>>();
let mut gossip_engine = GossipEngine::<Block>::new(
network.clone(),
sync.clone(),
notification_service,
protocol.clone(),
Arc::new(TestValidator {}),
None,
);
let (txs, mut rxs) = channels
.iter()
.map(|ChannelLengthAndTopic { length, topic }| (*topic, channel(*length)))
.fold((vec![], vec![]), |mut acc, (topic, (tx, rx))| {
acc.0.push((topic, tx));
acc.1.push((topic, rx));
acc
});
for (topic, tx) in txs {
match gossip_engine.message_sinks.get_mut(&topic) {
Some(entry) => entry.push(tx),
None => {
gossip_engine.message_sinks.insert(topic, vec![tx]);
},
}
}
tx.start_send(NotificationEvent::NotificationStreamOpened {
peer: remote_peer,
direction: Direction::Inbound,
negotiated_fallback: None,
handshake: Roles::FULL.encode(),
})
.unwrap();
for (i_notification, messages) in notifications.iter().enumerate() {
let messages: Vec<Vec<u8>> = messages
.into_iter()
.enumerate()
.map(|(i_message, Message { topic })| {
let mut message = topic.as_bytes().to_vec();
message.push(i_notification.try_into().unwrap());
message.push(i_message.try_into().unwrap());
message.into()
})
.collect();
for message in messages {
tx.start_send(NotificationEvent::NotificationReceived {
peer: remote_peer,
notification: message,
})
.unwrap();
}
}
let mut received_msgs_per_topic_all_chan = HashMap::<H256, _>::new();
block_on(poll_fn(|cx| {
loop {
if let Poll::Ready(()) = gossip_engine.poll_unpin(cx) {
unreachable!(
"Event stream sender side is not dropped, thus gossip engine does not \
terminate",
);
}
let mut progress = false;
for (topic, rx) in rxs.iter_mut() {
match rx.poll_next_unpin(cx) {
Poll::Ready(Some(_)) => {
progress = true;
received_msgs_per_topic_all_chan
.entry(*topic)
.and_modify(|e| *e += 1)
.or_insert(1);
},
Poll::Ready(None) => {
unreachable!("Sender side of channel is never dropped")
},
Poll::Pending => {},
}
}
if !progress {
break
}
}
Poll::Ready(())
}));
for (expected_topic, expected_num) in expected_msgs_per_topic_all_chan.iter() {
assert_eq!(
received_msgs_per_topic_all_chan.get(&expected_topic).unwrap_or(&0),
expected_num,
);
}
for (received_topic, received_num) in expected_msgs_per_topic_all_chan.iter() {
assert_eq!(
expected_msgs_per_topic_all_chan.get(&received_topic).unwrap_or(&0),
received_num,
);
}
}
prop(vec![], vec![vec![Message { topic: H256::default() }]]);
prop(
vec![ChannelLengthAndTopic { length: 71, topic: H256::default() }],
vec![vec![Message { topic: H256::default() }]],
);
QuickCheck::new().quickcheck(prop as fn(_, _))
}
}