Add toggle to disable accepting new connections

src/connections.rs

@@ -6,12 +6,14 @@ use std::{
     time::{Duration, Instant},
 };
 
+use serde::{Serialize, de::DeserializeOwned};
 use thiserror::Error;
 
 use crate::{
     PeerConfig, PeerMessage,
-    listener::{DATAGRAM_SIZE, Datagram, ListenerError},
-    package::Package,
+    listener::{DATAGRAM_SIZE, ListenerError},
+    package::{Message, Messages, Package},
+    stats::NetStats,
 };
 
 /// Represents an error during sending a package for any number of reason
@@ -39,22 +41,29 @@ pub enum ConnectionError {
 }
 
 /// Manages connections for the Peer
-pub(crate) struct ConnectionManager {
+pub(crate) struct ConnectionManager<
+    T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + 'static,
+> {
     /// Wraps UdpSocket with some helper methods
     pub udp: UdpWrapper,
-    connections: HashMap<SocketAddr, Connection>,
+    connections: HashMap<SocketAddr, Connection<T>>,
     closing_since: Option<Instant>,
     config: PeerConfig,
+
+    accepting_connections: bool,
 }
 
-impl ConnectionManager {
+impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + 'static>
+    ConnectionManager<T>
+{
     /// Creates a new ConnectionManager with the given socket and config
-    pub fn new(socket: Arc<UdpSocket>, config: PeerConfig) -> ConnectionManager {
+    pub fn new(socket: Arc<UdpSocket>, config: PeerConfig) -> ConnectionManager<T> {
         ConnectionManager {
             udp: UdpWrapper::new(socket),
             connections: HashMap::new(),
             closing_since: None,
             config,
+            accepting_connections: true,
         }
     }
 
@@ -94,8 +103,9 @@ impl ConnectionManager {
                 if duration > self.config.ping_interval {
                     match conn.state {
                         ConnectionState::ReceivingConnection | ConnectionState::Connecting => {
-                            match self.udp.send_to(*addr, Package::Hello) {
-                                Ok(_) => {
+                            match self.udp.send_to(*addr, Package::<T>::Hello) {
+                                Ok(bytes) => {
+                                    conn.bytes_tx += bytes;
                                     conn.last_sent_ping = now;
                                 }
                                 Err(err) => {
@@ -106,8 +116,9 @@ impl ConnectionManager {
                             }
                         }
                         ConnectionState::Connected | ConnectionState::ConnectingNearlyReady => {
-                            match self.udp.send_to(*addr, Package::Ping) {
-                                Ok(_) => {
+                            match self.udp.send_to(*addr, Package::<T>::Ping) {
+                                Ok(bytes) => {
+                                    conn.bytes_tx += bytes;
                                     conn.last_sent_ping = now;
                                 }
                                 Err(err) => {
@@ -118,8 +129,9 @@ impl ConnectionManager {
                             }
                         }
                         ConnectionState::Closing | ConnectionState::ReceivingClosing => {
-                            match self.udp.send_to(*addr, Package::Close) {
-                                Ok(_) => {
+                            match self.udp.send_to(*addr, Package::<T>::Close) {
+                                Ok(bytes) => {
+                                    conn.bytes_tx += bytes;
                                     conn.last_sent_ping = now;
                                 }
                                 Err(err) => {
@@ -136,6 +148,77 @@ impl ConnectionManager {
         }
     }
 
+    pub fn send_queued_messages(&mut self) {
+        let now = Instant::now();
+        for (addr, conn) in &mut self.connections {
+            if now - conn.last_msg_send > self.config.message_retry {
+                let messages = Messages {
+                    ack: conn.ack,
+                    reliable: conn.reliable_queue.clone(),
+                    unreliable: Vec::new(),
+                    messages_sent: conn.messages_sent + 1,
+                };
+                match self.udp.send_to(*addr, Package::Messages(messages)) {
+                    Ok(bytes) => {
+                        conn.bytes_tx += bytes;
+                        conn.last_msg_send = now;
+                        conn.messages_sent += 1;
+                    }
+                    Err(err) => {
+                        conn.closing_since = Instant::now();
+                        conn.state = ConnectionState::Error;
+                        conn.error = Some(ConnectionError::SendError(err));
+                    }
+                }
+            }
+        }
+    }
+
+    pub fn send(&mut self, addr: &SocketAddr, message: T, reliable: bool) {
+        if let Some(conn) = self.connections.get_mut(addr) {
+            if conn.state == ConnectionState::Connected {
+                let message_id = conn.message_counter;
+                if reliable {
+                    // Only increment message counter for reliable messages
+                    conn.message_counter += 1;
+                }
+
+                let message = Message {
+                    message_id,
+                    message,
+                };
+
+                let mut reliables = conn.reliable_queue.clone();
+                let mut unreliables = Vec::new();
+                if reliable {
+                    reliables.push(message.clone());
+                    conn.reliable_queue.push(message);
+                } else {
+                    unreliables.push(message);
+                }
+
+                let messages = Messages {
+                    ack: conn.ack,
+                    reliable: reliables,
+                    unreliable: unreliables,
+                    messages_sent: conn.messages_sent + 1,
+                };
+                match self.udp.send_to(*addr, Package::Messages(messages)) {
+                    Ok(bytes) => {
+                        conn.bytes_tx += bytes;
+                        conn.last_msg_send = Instant::now();
+                        conn.messages_sent += 1;
+                    }
+                    Err(err) => {
+                        conn.closing_since = Instant::now();
+                        conn.state = ConnectionState::Error;
+                        conn.error = Some(ConnectionError::SendError(err));
+                    }
+                }
+            }
+        }
+    }
+
     /// Close the connection to another peer at the given address
     pub fn close_connection(&mut self, addr: &SocketAddr) {
         if let Some(conn) = self.connections.get_mut(addr) {
@@ -159,7 +242,16 @@ impl ConnectionManager {
 
     /// Handle Package from a remote peer
     #[must_use]
-    pub fn handle_package(&mut self, package: Package, addr: &SocketAddr) -> Vec<PeerMessage> {
+    pub fn handle_package(
+        &mut self,
+        package: Package<T>,
+        addr: &SocketAddr,
+        bytes: usize,
+    ) -> Vec<PeerMessage<T>> {
+        if let Some(conn) = self.connections.get_mut(addr) {
+            conn.bytes_rx += bytes;
+        }
+
         let mut messages = Vec::new();
 
         match package {
@@ -169,10 +261,15 @@ impl ConnectionManager {
                         conn.state = ConnectionState::ConnectingNearlyReady;
                     }
                 } else {
-                    self.connections.insert(
-                        *addr,
-                        Connection::from(*addr, ConnectionState::ReceivingConnection),
-                    );
+                    if !self.accepting_connections {
+                        self.connections
+                            .insert(*addr, Connection::from(*addr, ConnectionState::Closing));
+                    } else {
+                        self.connections.insert(
+                            *addr,
+                            Connection::from(*addr, ConnectionState::ReceivingConnection),
+                        );
+                    }
                 }
             }
             Package::Ping => {
@@ -186,8 +283,10 @@ impl ConnectionManager {
 
                     if conn.state == ConnectionState::Connected {
                         conn.last_recv_ping = Instant::now();
-                        match self.udp.send_to(*addr, Package::Pong) {
-                            Ok(_) => {}
+                        match self.udp.send_to(*addr, Package::<T>::Pong) {
+                            Ok(bytes) => {
+                                conn.bytes_tx += bytes;
+                            }
                             Err(err) => {
                                 conn.closing_since = Instant::now();
                                 conn.state = ConnectionState::Error;
@@ -230,6 +329,51 @@ impl ConnectionManager {
                     self.connections.remove(addr);
                 }
             }
+            Package::Messages(msgs) => {
+                if let Some(conn) = self.connections.get_mut(addr) {
+                    if conn.state == ConnectionState::Connected {
+                        // Store the expected number of messages to have received
+                        // from the given connection
+                        conn.messages_expected = msgs.messages_sent.max(conn.messages_expected);
+                        conn.messages_received += 1;
+
+                        // Clear all messages from the queue that the other peer has
+                        // already seen
+                        conn.reliable_queue.retain(|v| v.message_id > msgs.ack);
+
+                        // Process all messages that have not been previously seen
+                        // before, that is to say the connection's ack value is
+                        // lower than the message's id
+                        for reliable_msg in &msgs.reliable {
+                            if reliable_msg.message_id > conn.ack {
+                                messages.push(PeerMessage::Message(
+                                    conn.clone(),
+                                    reliable_msg.message.clone(),
+                                ));
+                            }
+                        }
+
+                        // Set connection's ack value to be at least the highest
+                        // message_id that was received, so that we inform the
+                        // remote peer not to send these messages again.
+                        conn.ack = msgs
+                            .reliable
+                            .iter()
+                            .map(|r| r.message_id)
+                            .max()
+                            .unwrap_or(0)
+                            .max(conn.ack);
+
+                        // Always process all unreliable messages
+                        for unreliable_msg in &msgs.unreliable {
+                            messages.push(PeerMessage::Message(
+                                conn.clone(),
+                                unreliable_msg.message.clone(),
+                            ));
+                        }
+                    }
+                }
+            }
         }
 
         messages
@@ -247,7 +391,7 @@ impl ConnectionManager {
     }
 
     /// Remove any connections that have closed or errored.
-    pub fn purge_old_connections(&mut self) -> Vec<PeerMessage> {
+    pub fn purge_old_connections(&mut self) -> Vec<PeerMessage<T>> {
         let mut messages = Vec::new();
 
         let now = Instant::now();
@@ -273,6 +417,14 @@ impl ConnectionManager {
 
         messages
     }
+
+    pub fn connections(&self) -> Vec<Connection<T>> {
+        self.connections.values().cloned().collect()
+    }
+
+    pub fn set_accepting_connections(&mut self, accepting: bool) {
+        self.accepting_connections = accepting;
+    }
 }
 
 /// Wrap UdpSocket with some helper methods
@@ -298,7 +450,11 @@ impl UdpWrapper {
     }
 
     /// Send a Package to a remote peer at addr
-    fn send_to(&self, addr: SocketAddr, package: Package) -> Result<(), SendError> {
+    fn send_to<T: Serialize>(
+        &self,
+        addr: SocketAddr,
+        package: Package<T>,
+    ) -> Result<usize, SendError> {
         if let Some(socket) = &self.socket {
             let mut buf = Vec::new();
             let cursor = Cursor::new(&mut buf);
@@ -310,7 +466,7 @@ impl UdpWrapper {
             socket
                 .send_to(&buf, addr)
                 .map_err(|e| SendError::SendError(e))
-                .map(|_| ())
+                .map(|_| buf.len())
         } else {
             Err(SendError::SocketDisconnected)
         }
@@ -318,7 +474,7 @@ impl UdpWrapper {
 }
 
 /// Represents a single connection from the peer to another peer
-pub struct Connection {
+pub struct Connection<T: Clone> {
     /// Address of the connection
     pub address: SocketAddr,
     /// "ping", as in the time between the last sent ping and the last received
@@ -330,12 +486,24 @@ pub struct Connection {
     last_recv_close: Instant,
     closing_since: Instant,
 
+    message_counter: u64,
+    ack: u64,
+    reliable_queue: Vec<Message<T>>,
+    last_msg_send: Instant,
+
     /// What state is the connection currently in?
     pub state: ConnectionState,
     error: Option<ConnectionError>,
+
+    // Stats
+    messages_sent: usize,
+    messages_received: usize,
+    messages_expected: usize,
+    bytes_tx: usize,
+    bytes_rx: usize,
 }
 
-impl Clone for Connection {
+impl<T: Clone> Clone for Connection<T> {
     fn clone(&self) -> Self {
         Self {
             address: self.address.clone(),
@@ -344,14 +512,27 @@ impl Clone for Connection {
             last_recv_ping: self.last_recv_ping.clone(),
             last_recv_close: self.last_recv_close.clone(),
             closing_since: self.closing_since.clone(),
+
+            message_counter: self.message_counter,
+            ack: self.ack,
+            reliable_queue: self.reliable_queue.clone(),
+            last_msg_send: self.last_msg_send.clone(),
+
             state: self.state.clone(),
             error: None,
+
+            // Stats
+            messages_sent: self.messages_sent,
+            messages_received: self.messages_received,
+            messages_expected: self.messages_expected,
+            bytes_tx: self.bytes_tx,
+            bytes_rx: self.bytes_rx,
         }
     }
 }
 
-impl Connection {
-    pub fn from(address: SocketAddr, state: ConnectionState) -> Connection {
+impl<T: Clone> Connection<T> {
+    pub fn from(address: SocketAddr, state: ConnectionState) -> Connection<T> {
         Connection {
             address,
             ping: Duration::default(),
@@ -359,8 +540,29 @@ impl Connection {
             last_recv_ping: Instant::now(),
             last_recv_close: Instant::now(),
             closing_since: Instant::now(),
+
+            message_counter: 1,
+            ack: 0,
+            reliable_queue: Vec::new(),
+            last_msg_send: Instant::now() - Duration::from_hours(1),
+
             state,
             error: None,
+
+            messages_sent: 0,
+            messages_received: 0,
+            messages_expected: 0,
+            bytes_tx: 0,
+            bytes_rx: 0,
+        }
+    }
+
+    pub fn statistics(&self) -> NetStats {
+        NetStats {
+            bytes_tx: self.bytes_tx,
+            bytes_rx: self.bytes_rx,
+            messages_received: self.messages_received,
+            messages_expected: self.messages_expected,
         }
     }
 }

src/lib.rs

@@ -10,16 +10,19 @@ use std::{
     time::Duration,
 };
 
+use serde::{Serialize, de::DeserializeOwned};
 use thiserror::*;
 
 use crate::{
     connections::{Connection, ConnectionError, ConnectionManager},
     listener::{Listener, ListenerMessage},
+    stats::NetStats,
 };
 
 pub mod connections;
 pub(crate) mod listener;
 pub(crate) mod package;
+pub mod stats;
 
 /// Error for the Peer
 #[derive(Error, Debug)]
@@ -31,13 +34,15 @@ pub enum PeerError {
 }
 
 /// Possible messages from occurring events from the peer
-pub enum PeerMessage {
+pub enum PeerMessage<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> {
     /// A new connection has been connected
-    NewConnection(Connection),
+    NewConnection(Connection<T>),
     /// An existing connection has disconnected, with an optional error
-    Disconnected(Connection, Option<ConnectionError>),
+    Disconnected(Connection<T>, Option<ConnectionError>),
     /// The Peer has been closed
     Closed,
+    /// A single message of type T
+    Message(Connection<T>, T),
 }
 
 /// Optional configuration available for peers, mainly configuration of
@@ -47,6 +52,7 @@ pub struct PeerConfig {
     ping_interval: Duration,
     timeout: Duration,
     disconnect_timeout: Duration,
+    message_retry: Duration,
 }
 
 impl Default for PeerConfig {
@@ -55,6 +61,7 @@ impl Default for PeerConfig {
             ping_interval: Duration::from_millis(100),
             timeout: Duration::from_millis(2000),
             disconnect_timeout: Duration::from_millis(500),
+            message_retry: Duration::from_millis(100),
         }
     }
 }
@@ -83,19 +90,29 @@ impl PeerConfig {
             ..self
         }
     }
+
+    /// Sets the duration which is awaited before any queued reliable messages
+    /// are sent again. This timer is also reset every time any new messages are
+    /// sent manually.
+    pub fn with_retry(self, retry: Duration) -> PeerConfig {
+        PeerConfig {
+            message_retry: retry,
+            ..self
+        }
+    }
 }
 
-pub struct Peer {
-    connection_mgr: ConnectionManager,
+pub struct Peer<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + 'static> {
+    connection_mgr: ConnectionManager<T>,
     closed: Arc<AtomicBool>,
-    receiver: Receiver<ListenerMessage>,
-    messages: VecDeque<PeerMessage>,
+    receiver: Receiver<ListenerMessage<T>>,
+    messages: VecDeque<PeerMessage<T>>,
 }
 
-impl Peer {
+impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Peer<T> {
     /// Bind given port to listen to incoming messages. Creates a new Peer that
     /// can establish new connections.
-    pub fn listen(port: Option<u16>, config: PeerConfig) -> Result<Peer, PeerError> {
+    pub fn listen(port: Option<u16>, config: PeerConfig) -> Result<Peer<T>, PeerError> {
         let socket = Arc::new(
             UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], port.unwrap_or(0))))
                 .map_err(|e| PeerError::BindError(e))?,
@@ -133,7 +150,7 @@ impl Peer {
     /// possible events, such as connections and messages. Should be called as
     /// often as possible.
     #[must_use]
-    pub fn poll(&mut self) -> Result<Option<PeerMessage>, PeerError> {
+    pub fn poll(&mut self) -> Result<Option<PeerMessage<T>>, PeerError> {
         if self.connection_mgr.udp.is_closed() {
             if self.closed.load(Ordering::Acquire) {
                 self.messages.push_back(PeerMessage::Closed);
@@ -142,6 +159,7 @@ impl Peer {
         }
 
         self.connection_mgr.send_pings();
+        self.connection_mgr.send_queued_messages();
 
         let mut msg;
         while {
@@ -153,9 +171,12 @@ impl Peer {
             msg.is_some()
         } {
             match msg.unwrap() {
-                ListenerMessage::Package(package, socket_addr) => {
-                    self.messages
-                        .extend(self.connection_mgr.handle_package(package, &socket_addr));
+                ListenerMessage::Package(package, socket_addr, bytes) => {
+                    self.messages.extend(self.connection_mgr.handle_package(
+                        package,
+                        &socket_addr,
+                        bytes,
+                    ));
                 }
                 ListenerMessage::PackageError(listener_error, socket_addr) => {
                     println!("Error: {}", listener_error);
@@ -180,8 +201,57 @@ impl Peer {
         Ok(self.messages.pop_front())
     }
 
+    /// Send a reliable message, which is retried until the remote peer
+    /// acknowledges it.
+    pub fn send_reliable(&mut self, to: &SocketAddr, message: T) {
+        self.connection_mgr.send(to, message, true);
+    }
+
+    /// Sends an unreliable message, which is never retried again.
+    pub fn send_unreliable(&mut self, to: &SocketAddr, message: T) {
+        self.connection_mgr.send(to, message, false);
+    }
+
+    /// Send a reliable message to every connection currently connected
+    pub fn broadcast_reliable(&mut self, message: T) {
+        for connection in self.connections() {
+            self.connection_mgr
+                .send(&connection.address, message.clone(), true);
+        }
+    }
+
+    /// Send an unreliable message to every connection currently connected
+    pub fn broadcast_unreliable(&mut self, message: T) {
+        for connection in self.connections() {
+            self.connection_mgr
+                .send(&connection.address, message.clone(), false);
+        }
+    }
+
     /// Close the Peer, disconnecting all connections.
     pub fn close(&mut self) {
         self.connection_mgr.close();
     }
+
+    /// List of all current connections
+    pub fn connections(&self) -> Vec<Connection<T>> {
+        self.connection_mgr.connections()
+    }
+
+    pub fn set_accepting_connections(&mut self, accepting: bool) {
+        self.connection_mgr.set_accepting_connections(accepting);
+    }
+
+    pub fn statistics(&self) -> NetStats {
+        let conns = self.connection_mgr.connections();
+        let all_stats = conns.iter().map(|c| c.statistics());
+        let mut total_stats = NetStats::default();
+        for stats in all_stats {
+            total_stats.bytes_rx += stats.bytes_rx;
+            total_stats.bytes_tx += stats.bytes_tx;
+            total_stats.messages_expected += stats.messages_expected;
+            total_stats.messages_received += stats.messages_received;
+        }
+        total_stats
+    }
 }

src/listener.rs

@@ -1,9 +1,11 @@
 use std::{
     io::Cursor,
+    marker::PhantomData,
     net::{SocketAddr, UdpSocket},
     sync::{Arc, mpsc::Sender},
 };
 
+use serde::{Serialize, de::DeserializeOwned};
 use thiserror::Error;
 
 use crate::package::Package;
@@ -12,14 +14,19 @@ pub const DATAGRAM_SIZE: usize = 102400;
 pub type Datagram = [u8; DATAGRAM_SIZE];
 
 /// Listener thread for Peer's UdpSocket
-pub struct Listener {
+pub struct Listener<T: Clone + Serialize + DeserializeOwned> {
     socket: Arc<UdpSocket>,
-    sender: Sender<ListenerMessage>,
+    sender: Sender<ListenerMessage<T>>,
+    pd: PhantomData<T>,
 }
 
-impl Listener {
-    pub fn new(socket: Arc<UdpSocket>, sender: Sender<ListenerMessage>) -> Listener {
-        Listener { socket, sender }
+impl<'a, T: Clone + Serialize + DeserializeOwned> Listener<T> {
+    pub fn new(socket: Arc<UdpSocket>, sender: Sender<ListenerMessage<T>>) -> Listener<T> {
+        Listener {
+            socket,
+            sender,
+            pd: PhantomData::default(),
+        }
     }
 
     /// Read from the socket, parse packages and emit messages.
@@ -29,9 +36,12 @@ impl Listener {
             Ok((num_bytes, from_addr)) => {
                 let mut ciborium_buf: Datagram = [0; _];
                 let bytes = Cursor::new(&mut buffer[..num_bytes]);
-                let res = ciborium::from_reader_with_buffer::<Package, _>(bytes, &mut ciborium_buf);
+                let res =
+                    ciborium::from_reader_with_buffer::<Package<T>, _>(bytes, &mut ciborium_buf);
                 match res {
-                    Ok(pkg) => self.sender.send(ListenerMessage::Package(pkg, from_addr)),
+                    Ok(pkg) => self
+                        .sender
+                        .send(ListenerMessage::Package(pkg, from_addr, num_bytes)),
                     Err(_) => self.sender.send(ListenerMessage::PackageError(
                         ListenerError::ParseError,
                         from_addr,
@@ -52,8 +62,8 @@ pub enum ListenerError {
     ParseError,
 }
 
-pub enum ListenerMessage {
-    Package(Package, SocketAddr),
+pub enum ListenerMessage<T: Clone + Serialize + DeserializeOwned> {
+    Package(Package<T>, SocketAddr, usize),
     PackageError(ListenerError, SocketAddr),
     Error(std::io::Error),
 }

src/package.rs

@@ -1,9 +1,24 @@
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Serialize, Deserialize)]
-pub enum Package {
+pub enum Package<T> {
     Hello,
     Ping,
     Pong,
     Close,
+    Messages(Messages<T>),
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Messages<T> {
+    pub ack: u64,
+    pub messages_sent: usize,
+    pub reliable: Vec<Message<T>>,
+    pub unreliable: Vec<Message<T>>,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct Message<T> {
+    pub message_id: u64,
+    pub message: T,
 }

src/stats.rs

@@ -0,0 +1,7 @@
+#[derive(Default)]
+pub struct NetStats {
+    pub bytes_tx: usize,
+    pub bytes_rx: usize,
+    pub messages_received: usize,
+    pub messages_expected: usize,
+}