Update implementation

statement-gossip-subsystem/Cargo.toml

@@ -11,8 +11,5 @@ polkadot-network = { path = "../network" }
 statement-table = { package = "polkadot-statement-table", path = "../statement-table" }
 primitives = { package = "polkadot-primitives", path = "../primitives" }
 overseer = { path = "../overseer" }
-async-std = "1.5.0"
 exit-future = "0.2.0"
 futures = "0.3.5"
-sc-service = { git = "https://github.com/paritytech/substrate", branch = "master" }
-env_logger = "0.7.1"

statement-gossip-subsystem/src/lib.rs

@@ -4,87 +4,66 @@ use std::collections::HashMap;
 use futures::channel::mpsc::{Sender, Receiver, channel};
 use futures::prelude::*;
 use futures::future::ready;
-use futures::task::{Spawn, SpawnExt};
 use polkadot_network::legacy::GossipService;
 use polkadot_network::legacy::gossip::{GossipMessage, GossipStatement};
 use overseer::*;
 
-#[derive(Debug, Clone)]
-pub enum StatementGossipMessage {
-    ToGossip { relay_parent: Hash, statement: SignedStatement },
-    Received { relay_parent: Hash, statement: SignedStatement }
-}
-
 type Jobs = HashMap<Hash, (exit_future::Signal, Sender<SignedStatement>)>;
+type Message = StatementGossipSubsystemMessage;
 
-#[derive(Debug, Clone)]
-enum Message {
-    StartWork(Hash),
-    StopWork(Hash),
-    StatementGossip(StatementGossipMessage),
-}
-
-pub struct StatementGossipSubsystem<GS, SP> {
+pub struct StatementGossipSubsystem<GS> {
     /// This comes from the legacy networking code, so it is likely to be changed.
     gossip_service: GS,
-    /// A `futures` Spawner such as `sc_service::SpawnTaskHandle`.
-    spawner: SP
 }
 
-impl<GS, SP> StatementGossipSubsystem<GS, SP> {
-    pub fn new(gossip_service: GS, spawner: SP) -> Self {
+impl<GS> StatementGossipSubsystem<GS> {
+    pub fn new(gossip_service: GS) -> Self {
         Self {
-            gossip_service, spawner
+            gossip_service
         }
     }
 }
 
-impl<GS: GossipService + Clone + Send + Sync + 'static, SP: Spawn + Clone + 'static> Subsystem<Message> for StatementGossipSubsystem<GS, SP> {
-    fn start(&mut self, rx: Receiver<Message>, mut tx: Sender<OverseerMessage<Message>>) -> SubsystemJob {
+impl<GS: GossipService + Clone + Send + Sync + 'static> Subsystem<Message> for StatementGossipSubsystem<GS> {
+    fn start(&mut self, mut ctx: SubsystemContext<Message>) -> SpawnedSubsystem {
         let mut jobs = Jobs::new();
         let gossip_service = self.gossip_service.clone();
-        let spawner = self.spawner.clone();
-        let (jobs_to_subsystem_s, jobs_to_subsystem_r) = channel(1024);
-
-        let incoming = rx.for_each(move |message| {
-            match message {
-                Message::StartWork(relay_parent) => {
-                    let (signal, exit) = exit_future::signal();
-                    let (subsystem_to_job_s, subsystem_to_job_r) = channel(1024);
-
-                    spawner.spawn(statement_gossip_job(
-                        gossip_service.clone(), relay_parent, exit, jobs_to_subsystem_s.clone(), subsystem_to_job_r,
-                    )).unwrap();
-
-                    jobs.insert(relay_parent, (signal, subsystem_to_job_s));
-                },
-                Message::StopWork(relay_parent) => {
-                    if let Some((signal, _)) = jobs.remove(&relay_parent) {
-                        signal.fire().unwrap();
-                    } else {
-                        println!("Error: No jobs for {:?} running", relay_parent);
-                    }
-                },
-                Message::StatementGossip(StatementGossipMessage::ToGossip { relay_parent, statement }) => {
-                    if let Some((_, sender)) = jobs.get_mut(&relay_parent) {
-                        sender.try_send(statement).unwrap();
-                    }
-                },
-                _ => {}
+        let (jobs_to_subsystem_s, mut jobs_to_subsystem_r) = channel(1024);
+
+        SpawnedSubsystem(Box::pin(async move {
+            loop {
+                match ctx.try_recv().await {
+                    Ok(Some(FromOverseer::Signal(OverseerSignal::StartWork(relay_parent)))) => {
+                        let (signal, exit) = exit_future::signal();
+                        let (subsystem_to_job_s, subsystem_to_job_r) = channel(1024);
+
+                        ctx.spawn(statement_gossip_job(
+                            gossip_service.clone(), relay_parent, exit, jobs_to_subsystem_s.clone(), subsystem_to_job_r,
+                        ).boxed()).await.unwrap();
+
+                        jobs.insert(relay_parent, (signal, subsystem_to_job_s));
+                    },
+                    Ok(Some(FromOverseer::Signal(OverseerSignal::StopWork(relay_parent)))) => {
+                        if let Some((signal, _)) = jobs.remove(&relay_parent) {
+                            signal.fire().unwrap();
+                        } else {
+                            println!("Error: No jobs for {:?} running", relay_parent);
+                        }
+                    },
+                    Ok(Some(FromOverseer::Communication { msg: Message::ToGossip { relay_parent, statement }})) => {
+                        if let Some((_, sender)) = jobs.get_mut(&relay_parent) {
+                            sender.try_send(statement).unwrap();
+                        }
+                    },
+                    Ok(Some(FromOverseer::Communication { msg: Message::Received { .. }})) => {},
+                    Ok(Some(FromOverseer::Signal(OverseerSignal::Conclude))) | Ok(None) | Err(_) => return,
+                }
+
+                if let Some(msg) = jobs_to_subsystem_r.next().await {
+                    let _ = ctx.send_msg(AllMessages::StatementGossip(msg)).await;
+                }
+
             }
-
-            futures::future::ready(())
-        });
-
-        let outgoing = jobs_to_subsystem_r.for_each(move |message| {
-            tx.try_send(OverseerMessage::SubsystemMessage { to: None, msg: Message::StatementGossip(message) });
-
-            futures::future::ready(())
-        });
-
-		SubsystemJob(Box::pin(async move {
-            futures::join!(incoming, outgoing);
-            Ok(())
         }))
 	}
 }
@@ -96,7 +75,7 @@ async fn statement_gossip_job<GS: GossipService>(
     // A future that resolves with the associated `exit_future::Signal` is fired.
     exit_future: exit_future::Exit,
     // A cloned sender of messages to the overseer. This channel is shared between all jobs.
-    mut overseer_sender: Sender<StatementGossipMessage>,
+    mut overseer_sender: Sender<StatementGossipSubsystemMessage>,
     // A receiver of messages from the subsystem. This channel is exclusive to this job.
     subsystem_receiver: Receiver<SignedStatement>,
 ) {
@@ -108,7 +87,7 @@ async fn statement_gossip_job<GS: GossipService>(
             _ => ready(None)
         })
         .for_each(|statement| {
-            overseer_sender.try_send(StatementGossipMessage::Received { relay_parent, statement: statement.signed_statement }).unwrap();
+            overseer_sender.try_send(StatementGossipSubsystemMessage::Received { relay_parent, statement: statement.signed_statement }).unwrap();
             ready(())
         })
         .fuse();
@@ -125,29 +104,3 @@ async fn statement_gossip_job<GS: GossipService>(
     }
     println!("Job for {:?} stopped", relay_parent);
 }
-
-#[test]
-fn subsystem() {
-    env_logger::init();
-
-    #[derive(Clone)]
-    struct AsyncStdSpawner;
-
-    impl Spawn for AsyncStdSpawner {
-        fn spawn_obj(&self, future: futures::task::FutureObj<'static, ()>) -> Result<(), futures::task::SpawnError> {
-            async_std::task::spawn(future);
-            Ok(())
-        }
-    }
-
-    let mock = polkadot_network::protocol::tests::MockGossip::default();
-
-    futures::executor::block_on(async {
-		let subsystems: Vec<Box<dyn Subsystem<Message>>> = vec![
-			Box::new(StatementGossipSubsystem::new(mock, AsyncStdSpawner)),
-		];
-
-		let overseer = Overseer::new(subsystems);
-		overseer.run().await;
-	});
-}