From 6794a683ce160ec49e96c69eab1c59d9687770d4 Mon Sep 17 00:00:00 2001 From: Sofia Date: Thu, 16 Jul 2026 05:19:23 +0300 Subject: [PATCH] Add NetStats --- src/connections.rs | 72 ++++++++++++++++++++++++++++++++++++++++------ src/lib.rs | 24 ++++++++++++++-- src/listener.rs | 6 ++-- src/package.rs | 1 + src/stats.rs | 7 +++++ 5 files changed, 96 insertions(+), 14 deletions(-) create mode 100644 src/stats.rs diff --git a/src/connections.rs b/src/connections.rs index dc74608..c69785c 100644 --- a/src/connections.rs +++ b/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 Connection match conn.state { ConnectionState::ReceivingConnection | ConnectionState::Connecting => { match self.udp.send_to(*addr, Package::::Hello) { - Ok(_) => { + Ok(bytes) => { + conn.bytes_tx += bytes; conn.last_sent_ping = now; } Err(err) => { @@ -109,7 +111,8 @@ impl Connection } ConnectionState::Connected | ConnectionState::ConnectingNearlyReady => { match self.udp.send_to(*addr, Package::::Ping) { - Ok(_) => { + Ok(bytes) => { + conn.bytes_tx += bytes; conn.last_sent_ping = now; } Err(err) => { @@ -121,7 +124,8 @@ impl Connection } ConnectionState::Closing | ConnectionState::ReceivingClosing => { match self.udp.send_to(*addr, Package::::Close) { - Ok(_) => { + Ok(bytes) => { + conn.bytes_tx += bytes; conn.last_sent_ping = now; } Err(err) => { @@ -146,10 +150,13 @@ impl 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 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 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 Connection &mut self, package: Package, addr: &SocketAddr, + bytes: usize, ) -> Vec> { + 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 Connection if conn.state == ConnectionState::Connected { conn.last_recv_ping = Instant::now(); match self.udp.send_to(*addr, Package::::Pong) { - Ok(_) => {} + Ok(bytes) => { + conn.bytes_tx += bytes; + } Err(err) => { conn.closing_since = Instant::now(); conn.state = ConnectionState::Error; @@ -300,6 +320,11 @@ impl 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, - ) -> Result<(), SendError> { + ) -> Result { 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 { /// What state is the connection currently in? pub state: ConnectionState, error: Option, + + // Stats + messages_sent: usize, + messages_received: usize, + messages_expected: usize, + bytes_tx: usize, + bytes_rx: usize, } impl Clone for Connection { @@ -465,6 +497,13 @@ impl Clone for Connection { 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 Connection { 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, } } } diff --git a/src/lib.rs b/src/lib.rs index 31b0e02..64348d5 100644 --- a/src/lib.rs +++ b/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 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); @@ -232,4 +237,17 @@ impl Peer { pub fn connections(&self) -> Vec> { 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 + } } diff --git a/src/listener.rs b/src/listener.rs index 57b1d27..17a62cb 100644 --- a/src/listener.rs +++ b/src/listener.rs @@ -39,7 +39,9 @@ impl<'a, T: Clone + Serialize + DeserializeOwned> Listener { let res = ciborium::from_reader_with_buffer::, _>(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 { - Package(Package, SocketAddr), + Package(Package, SocketAddr, usize), PackageError(ListenerError, SocketAddr), Error(std::io::Error), } diff --git a/src/package.rs b/src/package.rs index a908af9..386a62d 100644 --- a/src/package.rs +++ b/src/package.rs @@ -12,6 +12,7 @@ pub enum Package { #[derive(Debug, Serialize, Deserialize)] pub struct Messages { pub ack: u64, + pub messages_sent: usize, pub reliable: Vec>, pub unreliable: Vec>, } diff --git a/src/stats.rs b/src/stats.rs new file mode 100644 index 0000000..bdd9fc2 --- /dev/null +++ b/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, +}