Add NetStats

src/connections.rs

@@ -13,6 +13,7 @@ use crate::{
     PeerConfig, PeerMessage,
     listener::{DATAGRAM_SIZE, ListenerError},
     package::{Message, Messages, Package},
+    stats::NetStats,
 };
 
 /// Represents an error during sending a package for any number of reason
@@ -97,7 +98,8 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
                     match conn.state {
                         ConnectionState::ReceivingConnection | ConnectionState::Connecting => {
                             match self.udp.send_to(*addr, Package::<T>::Hello) {
-                                Ok(_) => {
+                                Ok(bytes) => {
+                                    conn.bytes_tx += bytes;
                                     conn.last_sent_ping = now;
                                 }
                                 Err(err) => {
@@ -109,7 +111,8 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
                         }
                         ConnectionState::Connected | ConnectionState::ConnectingNearlyReady => {
                             match self.udp.send_to(*addr, Package::<T>::Ping) {
-                                Ok(_) => {
+                                Ok(bytes) => {
+                                    conn.bytes_tx += bytes;
                                     conn.last_sent_ping = now;
                                 }
                                 Err(err) => {
@@ -121,7 +124,8 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
                         }
                         ConnectionState::Closing | ConnectionState::ReceivingClosing => {
                             match self.udp.send_to(*addr, Package::<T>::Close) {
-                                Ok(_) => {
+                                Ok(bytes) => {
+                                    conn.bytes_tx += bytes;
                                     conn.last_sent_ping = now;
                                 }
                                 Err(err) => {
@@ -146,10 +150,13 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
                     ack: conn.ack,
                     reliable: conn.reliable_queue.clone(),
                     unreliable: Vec::new(),
+                    messages_sent: conn.messages_sent,
                 };
                 match self.udp.send_to(*addr, Package::Messages(messages)) {
-                    Ok(_) => {
+                    Ok(bytes) => {
+                        conn.bytes_tx += bytes;
                         conn.last_msg_send = now;
+                        conn.messages_sent += 1;
                     }
                     Err(err) => {
                         conn.closing_since = Instant::now();
@@ -165,7 +172,10 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
         if let Some(conn) = self.connections.get_mut(addr) {
             if conn.state == ConnectionState::Connected {
                 let message_id = conn.message_counter;
-                conn.message_counter += 1;
+                if reliable {
+                    // Only increment message counter for reliable messages
+                    conn.message_counter += 1;
+                }
 
                 let message = Message {
                     message_id,
@@ -185,10 +195,13 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
                     ack: conn.ack,
                     reliable: reliables,
                     unreliable: unreliables,
+                    messages_sent: conn.messages_sent,
                 };
                 match self.udp.send_to(*addr, Package::Messages(messages)) {
-                    Ok(_) => {
+                    Ok(bytes) => {
+                        conn.bytes_tx += bytes;
                         conn.last_msg_send = Instant::now();
+                        conn.messages_sent += 1;
                     }
                     Err(err) => {
                         conn.closing_since = Instant::now();
@@ -227,7 +240,12 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
         &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 {
@@ -255,7 +273,9 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
                     if conn.state == ConnectionState::Connected {
                         conn.last_recv_ping = Instant::now();
                         match self.udp.send_to(*addr, Package::<T>::Pong) {
-                            Ok(_) => {}
+                            Ok(bytes) => {
+                                conn.bytes_tx += bytes;
+                            }
                             Err(err) => {
                                 conn.closing_since = Instant::now();
                                 conn.state = ConnectionState::Error;
@@ -300,6 +320,11 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Connection
             }
             Package::Messages(msgs) => {
                 if let Some(conn) = self.connections.get_mut(addr) {
+                    // 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);
@@ -406,7 +431,7 @@ impl UdpWrapper {
         &self,
         addr: SocketAddr,
         package: Package<T>,
-    ) -> Result<(), SendError> {
+    ) -> Result<usize, SendError> {
         if let Some(socket) = &self.socket {
             let mut buf = Vec::new();
             let cursor = Cursor::new(&mut buf);
@@ -418,7 +443,7 @@ impl UdpWrapper {
             socket
                 .send_to(&buf, addr)
                 .map_err(|e| SendError::SendError(e))
-                .map(|_| ())
+                .map(|_| buf.len())
         } else {
             Err(SendError::SocketDisconnected)
         }
@@ -446,6 +471,13 @@ pub struct Connection<T: Clone> {
     /// 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<T: Clone> Clone for Connection<T> {
@@ -465,6 +497,13 @@ impl<T: Clone> Clone for Connection<T> {
 
             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,
         }
     }
 }
@@ -486,6 +525,21 @@ impl<T: Clone> Connection<T> {
 
             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

@@ -16,11 +16,13 @@ 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)]
@@ -169,9 +171,12 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Peer<T> {
             msg.is_some()
         } {
             match msg.unwrap() {
-                ListenerMessage::Package(package, socket_addr) => {
+                ListenerMessage::Package(package, socket_addr, bytes) => {
-                    self.messages
+                    self.messages.extend(self.connection_mgr.handle_package(
-                        .extend(self.connection_mgr.handle_package(package, &socket_addr));
+                        package,
+                        &socket_addr,
+                        bytes,
+                    ));
                 }
                 ListenerMessage::PackageError(listener_error, socket_addr) => {
                     println!("Error: {}", listener_error);
@@ -232,4 +237,17 @@ impl<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static> Peer<T> {
     pub fn connections(&self) -> Vec<Connection<T>> {
         self.connection_mgr.connections()
     }
+
+    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

@@ -39,7 +39,9 @@ impl<'a, T: Clone + Serialize + DeserializeOwned> Listener<T> {
                 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,
@@ -61,7 +63,7 @@ pub enum ListenerError {
 }
 
 pub enum ListenerMessage<T: Clone + Serialize + DeserializeOwned> {
-    Package(Package<T>, SocketAddr),
+    Package(Package<T>, SocketAddr, usize),
     PackageError(ListenerError, SocketAddr),
     Error(std::io::Error),
 }

src/package.rs

@@ -12,6 +12,7 @@ pub enum Package<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>>,
 }

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,
+}