Browse Source

Acceptor class for inbound connections

narodnik 5 years ago
parent
commit
235099e440

+ 2 - 0
src/error.rs

@@ -37,6 +37,7 @@ pub enum Error {
     ConnectTimeout,
     ChannelStopped,
     ChannelTimeout,
+    ServiceStopped,
 }
 
 impl std::error::Error for Error {}
@@ -77,6 +78,7 @@ impl fmt::Display for Error {
             Error::ConnectTimeout => f.write_str("Connection timed out"),
             Error::ChannelStopped => f.write_str("Channel stopped"),
             Error::ChannelTimeout => f.write_str("Channel timed out"),
+            Error::ServiceStopped => f.write_str("Service stopped"),
         }
     }
 }

+ 88 - 0
src/net/acceptor.rs

@@ -0,0 +1,88 @@
+use futures::FutureExt;
+use log::*;
+use smol::{Async, Executor};
+use std::net::{SocketAddr, TcpListener};
+use std::sync::Arc;
+
+use crate::error::{Error, Result};
+use crate::net::{Channel, ChannelPtr, SettingsPtr};
+use crate::system::{StoppableTask, StoppableTaskPtr, Subscriber, SubscriberPtr, Subscription};
+
+pub type AcceptorPtr = Arc<Acceptor>;
+
+pub struct Acceptor {
+    channel_subscriber: SubscriberPtr<Result<ChannelPtr>>,
+    task: StoppableTaskPtr,
+    settings: SettingsPtr,
+}
+
+impl Acceptor {
+    pub fn new(settings: SettingsPtr) -> Arc<Self> {
+        Arc::new(Self {
+            channel_subscriber: Subscriber::new(),
+            task: StoppableTask::new(),
+            settings,
+        })
+    }
+
+    pub fn accept(
+        self: Arc<Self>,
+        accept_addr: SocketAddr,
+        executor: Arc<Executor<'_>>,
+    ) -> Result<()> {
+        let listener = Async::<TcpListener>::bind(accept_addr)?;
+        info!("Listening on {}", listener.get_ref().local_addr()?);
+
+        // Start detached task and return instantly
+        self.accept_or_stop(listener, executor);
+
+        Ok(())
+    }
+
+    pub async fn stop(&self) {
+        // Send stop signal
+        self.task.stop().await;
+    }
+
+    fn accept_or_stop(self: Arc<Self>, listener: Async<TcpListener>, executor: Arc<Executor<'_>>) {
+        self.task.clone().start(
+            self.clone().run_accept(listener),
+            |result| self.handle_stop(result),
+            executor,
+        );
+    }
+
+    async fn run_accept(self: Arc<Self>, listener: Async<TcpListener>) -> Result<()> {
+        loop {
+            match self.tick_accept(&listener).await {
+                Ok(channel) => {
+                    let channel_result = Arc::new(Ok(channel));
+                    self.channel_subscriber.notify(channel_result).await;
+                }
+                Err(err) => {
+                    error!("Error listening for connections: {}", err);
+                    return Err(Error::ServiceStopped);
+                }
+            }
+        }
+    }
+
+    async fn handle_stop(self: Arc<Self>, result: Result<()>) {
+        match result {
+            Ok(()) => panic!("Acceptor task should never complete without error status"),
+            Err(err) => {
+                // Send this error to all channel subscribers
+                let result = Arc::new(Err(err));
+                self.channel_subscriber.notify(result).await;
+            }
+        }
+    }
+
+    async fn tick_accept(&self, listener: &Async<TcpListener>) -> Result<ChannelPtr> {
+        let (stream, peer_addr) = listener.accept().await?;
+        info!("Accepted client: {}", peer_addr);
+
+        let channel = Channel::new(stream, peer_addr, self.settings.clone());
+        Ok(channel)
+    }
+}

+ 11 - 6
src/net/channel.rs

@@ -1,21 +1,23 @@
 use async_std::sync::Mutex;
-use std::sync::atomic::{AtomicBool, Ordering};
-use log::*;
-use futures::FutureExt;
 use futures::io::{ReadHalf, WriteHalf};
 use futures::AsyncReadExt;
+use futures::FutureExt;
+use log::*;
 use smol::{Async, Executor};
 use std::future::Future;
 use std::net::{SocketAddr, TcpStream};
 use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use crate::error::{Error, Result};
+use crate::net::message_subscriber::{
+    MessageSubscriber, MessageSubscriberPtr, MessageSubscription,
+};
 use crate::net::messages;
 use crate::net::settings::SettingsPtr;
-use crate::net::message_subscriber::{MessageSubscriberPtr, MessageSubscription, MessageSubscriber};
 use crate::net::utility::clone_net_error;
-use crate::system::{SubscriberPtr, Subscription, Subscriber};
+use crate::system::{Subscriber, SubscriberPtr, Subscription};
 
 pub type ChannelPtr = Arc<Channel>;
 
@@ -69,7 +71,10 @@ impl Channel {
         self.address
     }
 
-    pub async fn subscribe_msg(self: Arc<Self>, packet_type: messages::PacketType) -> MessageSubscription {
+    pub async fn subscribe_msg(
+        self: Arc<Self>,
+        packet_type: messages::PacketType,
+    ) -> MessageSubscription {
         self.message_subscriber.clone().subscribe(packet_type).await
     }
 

+ 9 - 6
src/net/hosts.rs

@@ -1,7 +1,7 @@
-use std::sync::Arc;
-use rand::seq::SliceRandom;
 use async_std::sync::Mutex;
+use rand::seq::SliceRandom;
 use std::net::SocketAddr;
+use std::sync::Arc;
 
 use crate::net::SettingsPtr;
 
@@ -9,14 +9,14 @@ pub type HostsPtr = Arc<Hosts>;
 
 pub struct Hosts {
     addrs: Mutex<Vec<SocketAddr>>,
-    settings: SettingsPtr
+    settings: SettingsPtr,
 }
 
 impl Hosts {
     pub fn new(settings: SettingsPtr) -> Arc<Self> {
         Arc::new(Self {
             addrs: Mutex::new(Vec::new()),
-            settings
+            settings,
         })
     }
 
@@ -25,7 +25,10 @@ impl Hosts {
     }
 
     pub async fn load(&self) -> Option<SocketAddr> {
-        self.addrs.lock().await.choose(&mut rand::thread_rng()).cloned()
+        self.addrs
+            .lock()
+            .await
+            .choose(&mut rand::thread_rng())
+            .cloned()
     }
 }
-

+ 16 - 21
src/net/message_subscriber.rs

@@ -1,7 +1,7 @@
-use std::collections::HashMap;
+use async_std::sync::Mutex;
 use rand::Rng;
+use std::collections::HashMap;
 use std::sync::Arc;
-use async_std::sync::Mutex;
 
 use crate::error::Result;
 use crate::net::messages::{Message, PacketType};
@@ -13,21 +13,16 @@ pub type MessageResult = Result<Arc<Message>>;
 pub type MessageSubscriptionID = u64;
 
 macro_rules! receive_message {
-    ($sub:expr, $message_type:path) => {
-        {
-            let wrapped_message = OwningRef::new($sub.receive().await?);
-
-            wrapped_message.map(|msg|
-                match msg {
-                    $message_type(msg_detail) => {
-                        msg_detail
-                    },
-                    _ => {
-                        panic!("Filter for receive sub invalid!");
-                    }
-            })
-        }
-    };
+    ($sub:expr, $message_type:path) => {{
+        let wrapped_message = OwningRef::new($sub.receive().await?);
+
+        wrapped_message.map(|msg| match msg {
+            $message_type(msg_detail) => msg_detail,
+            _ => {
+                panic!("Filter for receive sub invalid!");
+            }
+        })
+    }};
 }
 
 trait CloneMessageResult {
@@ -38,7 +33,7 @@ impl CloneMessageResult for Result<Arc<Message>> {
     fn clone(&self) -> Self {
         match self {
             Ok(message) => Ok(message.clone()),
-            Err(err) => Err(clone_net_error(err))
+            Err(err) => Err(clone_net_error(err)),
         }
     }
 }
@@ -47,7 +42,7 @@ pub struct MessageSubscription {
     id: MessageSubscriptionID,
     filter: PacketType,
     recv_queue: async_channel::Receiver<MessageResult>,
-    parent: Arc<MessageSubscriber>
+    parent: Arc<MessageSubscriber>,
 }
 
 impl MessageSubscription {
@@ -116,7 +111,7 @@ impl MessageSubscriber {
             id: sub_id,
             filter: packet_type,
             recv_queue: recvr,
-            parent: self.clone()
+            parent: self.clone(),
         }
     }
 
@@ -127,7 +122,7 @@ impl MessageSubscriber {
     pub async fn notify(&self, message_result: Result<Arc<Message>>) {
         for sub in (*self.subs.lock().await).values() {
             match sub.send(message_result.clone()).await {
-                Ok(()) => {},
+                Ok(()) => {}
                 Err(err) => {
                     panic!("Error returned sending message in notify() call! {}", err);
                 }

+ 10 - 20
src/net/messages.rs

@@ -180,26 +180,16 @@ impl Decodable for VerackMessage {
 impl Message {
     pub fn packet_type(&self) -> PacketType {
         match self {
-            Message::Ping => 
-                PacketType::Ping,
-            Message::Pong => 
-                PacketType::Pong,
-            Message::GetAddrs(message) => 
-                    PacketType::GetAddrs,
-            Message::Addrs(message) =>
-                    PacketType::Addrs,
-            Message::Sync => 
-                    PacketType::Sync,
-            Message::Inv(message) => 
-                    PacketType::Inv,
-            Message::GetSlabs(message) => 
-                    PacketType::GetSlabs,
-            Message::Slab(message) => 
-                    PacketType::Slab,
-            Message::Version(message) => 
-                    PacketType::Version,
-            Message::Verack(message) => 
-                    PacketType::Verack,
+            Message::Ping => PacketType::Ping,
+            Message::Pong => PacketType::Pong,
+            Message::GetAddrs(message) => PacketType::GetAddrs,
+            Message::Addrs(message) => PacketType::Addrs,
+            Message::Sync => PacketType::Sync,
+            Message::Inv(message) => PacketType::Inv,
+            Message::GetSlabs(message) => PacketType::GetSlabs,
+            Message::Slab(message) => PacketType::Slab,
+            Message::Version(message) => PacketType::Version,
+            Message::Verack(message) => PacketType::Verack,
         }
     }
 

+ 6 - 4
src/net/mod.rs

@@ -1,12 +1,13 @@
 use smol::Async;
 use std::net::TcpStream;
 
+pub mod acceptor;
 pub mod channel;
 pub mod connector;
 #[macro_use]
 pub mod message_subscriber;
-pub mod messages;
 pub mod hosts;
+pub mod messages;
 pub mod p2p;
 pub mod protocols;
 pub mod proxy;
@@ -16,10 +17,11 @@ pub mod utility;
 
 pub type AsyncTcpStream = async_dup::Arc<Async<TcpStream>>;
 
+pub use acceptor::{Acceptor, AcceptorPtr};
 pub use channel::{Channel, ChannelPtr};
 pub use connector::Connector;
-pub use message_subscriber::{MessageSubscription, MessageSubscriber};
-pub use hosts::{HostsPtr, Hosts};
+pub use hosts::{Hosts, HostsPtr};
+pub use message_subscriber::{MessageSubscriber, MessageSubscription};
 pub use p2p::P2p;
 pub use proxy::Proxy;
-pub use settings::{SettingsPtr, Settings};
+pub use settings::{Settings, SettingsPtr};

+ 2 - 2
src/net/p2p.rs

@@ -5,8 +5,8 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use crate::error::Result;
-use crate::net::sessions::{SeedSession, InboundSession};
-use crate::net::{Channel, ChannelPtr, Hosts, HostsPtr, Connector, Settings, SettingsPtr};
+use crate::net::sessions::{InboundSession, SeedSession};
+use crate::net::{Channel, ChannelPtr, Connector, Hosts, HostsPtr, Settings, SettingsPtr};
 
 pub type Pending<T> = Mutex<HashMap<SocketAddr, Arc<T>>>;
 

+ 2 - 3
src/net/protocols/mod.rs

@@ -1,8 +1,7 @@
-pub mod protocol_version;
 pub mod protocol_ping;
 pub mod protocol_seed;
+pub mod protocol_version;
 
-pub use protocol_version::ProtocolVersion;
 pub use protocol_ping::ProtocolPing;
 pub use protocol_seed::ProtocolSeed;
-
+pub use protocol_version::ProtocolVersion;

+ 7 - 4
src/net/protocols/protocol_ping.rs

@@ -1,11 +1,11 @@
-use rand::Rng;
 use futures::FutureExt;
+use rand::Rng;
 use smol::{Executor, Task};
 use std::sync::Arc;
 
 use crate::error::{Error, Result};
 use crate::net::messages;
-use crate::net::utility::{sleep, clone_net_error};
+use crate::net::utility::{clone_net_error, sleep};
 use crate::net::{ChannelPtr, SettingsPtr};
 
 pub struct ProtocolPing {
@@ -23,7 +23,11 @@ impl ProtocolPing {
     }
 
     async fn run_ping_pong(self: Arc<Self>) -> Result<()> {
-        let pong_sub = self.channel.clone().subscribe_msg(messages::PacketType::Pong).await;
+        let pong_sub = self
+            .channel
+            .clone()
+            .subscribe_msg(messages::PacketType::Pong)
+            .await;
 
         loop {
             // Wait channel_heartbeat amount of time
@@ -50,4 +54,3 @@ impl ProtocolPing {
         rng.gen()
     }
 }
-

+ 13 - 6
src/net/protocols/protocol_seed.rs

@@ -1,11 +1,11 @@
 use futures::FutureExt;
+use owning_ref::OwningRef;
 use smol::Executor;
 use std::sync::Arc;
-use owning_ref::OwningRef;
 
 use crate::error::{Error, Result};
-use crate::net::{ChannelPtr, HostsPtr, SettingsPtr};
 use crate::net::messages;
+use crate::net::{ChannelPtr, HostsPtr, SettingsPtr};
 
 pub struct ProtocolSeed {
     channel: ChannelPtr,
@@ -15,11 +15,19 @@ pub struct ProtocolSeed {
 
 impl ProtocolSeed {
     pub fn new(channel: ChannelPtr, hosts: HostsPtr, settings: SettingsPtr) -> Arc<Self> {
-        Arc::new(Self { channel, hosts, settings })
+        Arc::new(Self {
+            channel,
+            hosts,
+            settings,
+        })
     }
 
     pub async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) -> Result<()> {
-        let addr_sub = self.channel.clone().subscribe_msg(messages::PacketType::Addrs).await;
+        let addr_sub = self
+            .channel
+            .clone()
+            .subscribe_msg(messages::PacketType::Addrs)
+            .await;
 
         // Send own address to the seed server
         self.send_own_address().await?;
@@ -40,7 +48,7 @@ impl ProtocolSeed {
             Some(addr) => {
                 let addr = messages::Message::Addrs(messages::AddrsMessage { addrs: vec![addr] });
                 self.channel.clone().send(addr).await?;
-            },
+            }
             None => {
                 // Do nothing if external address is not configured
             }
@@ -48,4 +56,3 @@ impl ProtocolSeed {
         Ok(())
     }
 }
-

+ 6 - 2
src/net/protocols/protocol_version.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 
 use crate::error::{Error, Result};
 use crate::net::messages;
-use crate::net::utility::{sleep, clone_net_error};
+use crate::net::utility::{clone_net_error, sleep};
 use crate::net::{ChannelPtr, SettingsPtr};
 
 pub struct ProtocolVersion {
@@ -44,7 +44,11 @@ impl ProtocolVersion {
     }
 
     async fn recv_version(self: Arc<Self>) -> Result<()> {
-        let version_sub = self.channel.clone().subscribe_msg(messages::PacketType::Version).await;
+        let version_sub = self
+            .channel
+            .clone()
+            .subscribe_msg(messages::PacketType::Version)
+            .await;
 
         let version_msg = version_sub.receive().await?;
 

+ 42 - 5
src/net/sessions/inbound_session.rs

@@ -4,20 +4,58 @@ use std::net::SocketAddr;
 use std::sync::{Arc, Weak};
 
 use crate::error::{Error, Result};
-use crate::net::sessions::Session;
-use crate::net::{ChannelPtr, HostsPtr, Connector, P2p, SettingsPtr};
 use crate::net::protocols::{ProtocolPing, ProtocolSeed};
+use crate::net::sessions::Session;
+use crate::net::{Acceptor, AcceptorPtr};
+use crate::net::{ChannelPtr, Connector, HostsPtr, P2p, SettingsPtr};
 
 pub struct InboundSession {
-    p2p: Weak<P2p>
+    p2p: Weak<P2p>,
+    acceptor: AcceptorPtr,
 }
 
 impl InboundSession {
     pub fn new(p2p: Weak<P2p>) -> Arc<Self> {
-        Arc::new(Self { p2p })
+        let settings = {
+            let p2p = p2p.upgrade().unwrap();
+            p2p.settings()
+        };
+
+        let acceptor = Acceptor::new(settings);
+
+        Arc::new(Self { p2p, acceptor })
     }
 
     pub async fn start(self: Arc<Self>, executor: Arc<Executor<'_>>) -> Result<()> {
+        match self.p2p().settings().inbound {
+            Some(accept_addr) => {
+                self.start_accept_session(accept_addr, executor).await?;
+            }
+            None => {
+                info!("Not configured for accepting incoming connections.");
+            }
+        }
+
+        Ok(())
+    }
+
+    pub async fn stop(&self) {
+        self.acceptor.stop().await;
+    }
+
+    async fn start_accept_session(
+        self: Arc<Self>,
+        accept_addr: SocketAddr,
+        executor: Arc<Executor<'_>>,
+    ) -> Result<()> {
+        info!("Starting inbound session on {}", accept_addr);
+        match self.acceptor.clone().accept(accept_addr, executor) {
+            Ok(()) => {}
+            Err(err) => {
+                error!("Error starting listener: {}", err);
+                return Err(err);
+            }
+        }
         Ok(())
     }
 }
@@ -27,4 +65,3 @@ impl Session for InboundSession {
         self.p2p.upgrade().unwrap()
     }
 }
-

+ 25 - 8
src/net/sessions/seed_session.rs

@@ -4,12 +4,12 @@ use std::net::SocketAddr;
 use std::sync::{Arc, Weak};
 
 use crate::error::{Error, Result};
-use crate::net::sessions::Session;
-use crate::net::{ChannelPtr, HostsPtr, Connector, P2p, SettingsPtr};
 use crate::net::protocols::{ProtocolPing, ProtocolSeed};
+use crate::net::sessions::Session;
+use crate::net::{ChannelPtr, Connector, HostsPtr, P2p, SettingsPtr};
 
 pub struct SeedSession {
-    p2p: Weak<P2p>
+    p2p: Weak<P2p>,
 }
 
 impl SeedSession {
@@ -48,7 +48,11 @@ impl SeedSession {
         Ok(())
     }
 
-    async fn start_seed(self: Arc<Self>, seed: SocketAddr, executor: Arc<Executor<'_>>) -> Result<()> {
+    async fn start_seed(
+        self: Arc<Self>,
+        seed: SocketAddr,
+        executor: Arc<Executor<'_>>,
+    ) -> Result<()> {
         let (hosts, settings) = {
             let p2p = self.p2p.upgrade().unwrap();
             (p2p.hosts(), p2p.settings())
@@ -61,9 +65,12 @@ impl SeedSession {
 
                 info!("Connected seed [{}]", seed);
 
-                self.clone().register_channel(channel.clone(), executor.clone()).await?;
+                self.clone()
+                    .register_channel(channel.clone(), executor.clone())
+                    .await?;
 
-                self.attach_protocols(channel, hosts, settings, executor).await
+                self.attach_protocols(channel, hosts, settings, executor)
+                    .await
             }
             Err(err) => {
                 info!("Failure contacting seed [{}]: {}", seed, err);
@@ -72,7 +79,11 @@ impl SeedSession {
         }
     }
 
-    async fn register_channel(self: Arc<Self>, channel: ChannelPtr, executor: Arc<Executor<'_>>) -> Result<()> {
+    async fn register_channel(
+        self: Arc<Self>,
+        channel: ChannelPtr,
+        executor: Arc<Executor<'_>>,
+    ) -> Result<()> {
         let handshake_task = self.perform_handshake_protocols(channel.clone(), executor.clone());
 
         // start channel
@@ -81,7 +92,13 @@ impl SeedSession {
         handshake_task.await
     }
 
-    async fn attach_protocols(self: Arc<Self>, channel: ChannelPtr, hosts: HostsPtr, settings: SettingsPtr, executor: Arc<Executor<'_>>) -> Result<()> {
+    async fn attach_protocols(
+        self: Arc<Self>,
+        channel: ChannelPtr,
+        hosts: HostsPtr,
+        settings: SettingsPtr,
+        executor: Arc<Executor<'_>>,
+    ) -> Result<()> {
         let protocol_ping = ProtocolPing::new(channel.clone(), settings.clone());
         let ping_task = protocol_ping.start(executor.clone());
 

+ 6 - 2
src/net/sessions/session.rs

@@ -3,9 +3,9 @@ use smol::Executor;
 use std::sync::Arc;
 
 use crate::error::Result;
+use crate::net::p2p::P2pPtr;
 use crate::net::protocols::ProtocolVersion;
 use crate::net::ChannelPtr;
-use crate::net::p2p::P2pPtr;
 
 async fn remove_sub_on_stop(p2p: P2pPtr, channel: ChannelPtr) {
     // Subscribe to stop events
@@ -18,7 +18,11 @@ async fn remove_sub_on_stop(p2p: P2pPtr, channel: ChannelPtr) {
 
 #[async_trait]
 pub trait Session {
-    async fn perform_handshake_protocols(&self, channel: ChannelPtr, executor: Arc<Executor<'_>>) -> Result<()> {
+    async fn perform_handshake_protocols(
+        &self,
+        channel: ChannelPtr,
+        executor: Arc<Executor<'_>>,
+    ) -> Result<()> {
         let p2p = self.p2p();
 
         // Perform handshake

+ 1 - 1
src/net/settings.rs

@@ -1,5 +1,5 @@
-use std::sync::Arc;
 use std::net::SocketAddr;
+use std::sync::Arc;
 
 pub type SettingsPtr = Arc<Settings>;
 

+ 1 - 2
src/net/utility.rs

@@ -13,7 +13,6 @@ pub fn clone_net_error(error: &Error) -> Error {
         Error::ConnectTimeout => Error::ConnectTimeout,
         Error::ChannelStopped => Error::ChannelStopped,
         Error::ChannelTimeout => Error::ChannelTimeout,
-        _ => Error::OperationFailed
+        _ => Error::OperationFailed,
     }
 }
-

+ 5 - 0
src/system/mod.rs

@@ -0,0 +1,5 @@
+pub mod stoppable_task;
+pub mod subscriber;
+
+pub use stoppable_task::{StoppableTask, StoppableTaskPtr};
+pub use subscriber::{Subscriber, SubscriberPtr, Subscription};

+ 50 - 0
src/system/stoppable_task.rs

@@ -0,0 +1,50 @@
+use async_executor::Executor;
+use futures::Future;
+use futures::FutureExt;
+use std::sync::Arc;
+
+use crate::error::{Error, Result};
+
+pub type StoppableTaskPtr = Arc<StoppableTask>;
+
+pub struct StoppableTask {
+    stop_send: async_channel::Sender<()>,
+    stop_recv: async_channel::Receiver<()>,
+}
+
+impl StoppableTask {
+    pub fn new() -> Arc<Self> {
+        let (stop_send, stop_recv) = async_channel::unbounded();
+        Arc::new(Self {
+            stop_send,
+            stop_recv,
+        })
+    }
+
+    pub async fn stop(&self) {
+        // Ignore any errors from this send
+        let _ = self.stop_send.send(()).await;
+    }
+
+    pub fn start<'a, MainFut, StopFut, StopFn>(
+        self: Arc<Self>,
+        main: MainFut,
+        stop_handler: StopFn,
+        executor: Arc<Executor<'a>>,
+    ) where
+        MainFut: Future<Output = Result<()>> + Send + 'a,
+        StopFut: Future<Output = ()> + Send,
+        StopFn: FnOnce(Result<()>) -> StopFut + Send + 'a,
+    {
+        executor
+            .spawn(async move {
+                let result = futures::select! {
+                    _ = self.stop_recv.recv().fuse() => Err(Error::ServiceStopped),
+                    result = main.fuse() => result
+                };
+
+                stop_handler(result).await;
+            })
+            .detach();
+    }
+}

+ 79 - 0
src/system/subscriber.rs

@@ -0,0 +1,79 @@
+use async_std::sync::Mutex;
+use rand::Rng;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+pub type SubscriberPtr<T> = Arc<Subscriber<T>>;
+
+pub type SubscriptionID = u64;
+
+pub struct Subscription<T> {
+    id: SubscriptionID,
+    recv_queue: async_channel::Receiver<Arc<T>>,
+    parent: Arc<Subscriber<T>>,
+}
+
+impl<T> Subscription<T> {
+    pub async fn receive(&self) -> Arc<T> {
+        let message_result = self.recv_queue.recv().await;
+
+        match message_result {
+            Ok(message_result) => message_result,
+            Err(err) => {
+                panic!("MessageSubscription::receive() recv_queue failed! {}", err);
+            }
+        }
+    }
+
+    // Must be called manually since async Drop is not possible in Rust
+    pub async fn unsubscribe(&self) {
+        self.parent.clone().unsubscribe(self.id).await
+    }
+}
+
+// Simple broadcast (publish-subscribe) class
+pub struct Subscriber<T> {
+    subs: Mutex<HashMap<u64, async_channel::Sender<Arc<T>>>>,
+}
+
+impl<T> Subscriber<T> {
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {
+            subs: Mutex::new(HashMap::new()),
+        })
+    }
+
+    pub fn random_id() -> SubscriptionID {
+        let mut rng = rand::thread_rng();
+        rng.gen()
+    }
+
+    pub async fn subscribe(self: Arc<Self>) -> Subscription<T> {
+        let (sender, recvr) = async_channel::unbounded();
+
+        let sub_id = Self::random_id();
+
+        self.subs.lock().await.insert(sub_id, sender);
+
+        Subscription {
+            id: sub_id,
+            recv_queue: recvr,
+            parent: self.clone(),
+        }
+    }
+
+    async fn unsubscribe(self: Arc<Self>, sub_id: SubscriptionID) {
+        self.subs.lock().await.remove(&sub_id);
+    }
+
+    pub async fn notify(&self, message_result: Arc<T>) {
+        for sub in (*self.subs.lock().await).values() {
+            match sub.send(message_result.clone()).await {
+                Ok(()) => {}
+                Err(err) => {
+                    panic!("Error returned sending message in notify() call! {}", err);
+                }
+            }
+        }
+    }
+}