teascade/teanet-rust
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add close reasons

Browse Source
This commit is contained in:
Sofia 2026-08-04 20:48:45 +03:00
parent 22b08b78dd
commit 8837d1c6c3
3 changed files with 46 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
src/connections.rs
View File

@ -12,7 +12,7 @@ use thiserror::Error;
use crate::{ use crate::{
PeerConfig, PeerMessage, PeerConfig, PeerMessage,
listener::{DATAGRAM_SIZE, ListenerError}, listener::{DATAGRAM_SIZE, ListenerError},
package::{Message, Messages, Package}, package::{CloseReason, Message, Messages, Package},
stats::NetStats, stats::NetStats,
}; };
@ -71,7 +71,7 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
pub fn close(&mut self) { pub fn close(&mut self) {
self.closing_since = Some(Instant::now()); self.closing_since = Some(Instant::now());
for (addr, _) in self.connections.clone() { for (addr, _) in self.connections.clone() {
self.close_connection(&addr); self.close_connection(&addr, CloseReason::ShuttingDown);
} }
} }
@ -131,8 +131,9 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
} }
} }
} }
ConnectionState::Closing | ConnectionState::ReceivingClosing => { ConnectionState::Closing(reason)
match self.udp.send_to(*addr, Package::<T>::Close) { | ConnectionState::ReceivingClosing(reason) => {
match self.udp.send_to(*addr, Package::<T>::Close(reason)) {
Ok(bytes) => { Ok(bytes) => {
conn.bytes_tx += bytes; conn.bytes_tx += bytes;
conn.last_sent_ping = now; conn.last_sent_ping = now;
@ -223,11 +224,11 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
} }
/// Close the connection to another peer at the given address /// Close the connection to another peer at the given address
pub fn close_connection(&mut self, addr: &SocketAddr) { pub fn close_connection(&mut self, addr: &SocketAddr, reason: CloseReason) {
if let Some(conn) = self.connections.get_mut(addr) { if let Some(conn) = self.connections.get_mut(addr) {
if conn.state != ConnectionState::Closing { if !matches!(conn.state, ConnectionState::Closing(_)) {
conn.closing_since = Instant::now(); conn.closing_since = Instant::now();
conn.state = ConnectionState::Closing; conn.state = ConnectionState::Closing(reason);
} }
} }
} }
@ -265,8 +266,13 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
} }
} else { } else {
if !self.accepting_connections || self.config.identifier != identifier { if !self.accepting_connections || self.config.identifier != identifier {
self.connections self.connections.insert(
.insert(*addr, Connection::from(*addr, ConnectionState::Closing)); *addr,
Connection::from(
*addr,
ConnectionState::Closing(CloseReason::NotAcceptingConnections),
),
);
} else { } else {
self.connections.insert( self.connections.insert(
*addr, *addr,
@ -306,22 +312,25 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
} }
} }
} }
Package::Close => { Package::Close(reason) => {
let remove = if let Some(conn) = self.connections.get_mut(addr) { let remove = if let Some(conn) = self.connections.get_mut(addr) {
match conn.state { match conn.state {
ConnectionState::Closing => { ConnectionState::Closing(_) => {
messages messages.push(PeerMessage::Disconnected(
.push(PeerMessage::Disconnected(conn.clone(), conn.error.take())); conn.clone(),
conn.error.take(),
CloseReason::SelfClosed,
));
true true
} }
ConnectionState::Error => false, ConnectionState::Error => false,
ConnectionState::ReceivingClosing => { ConnectionState::ReceivingClosing(_) => {
conn.last_recv_close = Instant::now(); conn.last_recv_close = Instant::now();
false false
} }
_ => { _ => {
conn.last_recv_close = Instant::now(); conn.last_recv_close = Instant::now();
conn.state = ConnectionState::ReceivingClosing; conn.state = ConnectionState::ReceivingClosing(reason);
false false
} }
} }
@ -401,16 +410,20 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
self.connections.retain(|_, conn| { self.connections.retain(|_, conn| {
let retain = match conn.state { let retain = match conn.state {
ConnectionState::ReceivingClosing => { ConnectionState::ReceivingClosing(reason) => {
if (now - conn.last_recv_close) > self.config.disconnect_timeout { if (now - conn.last_recv_close) > self.config.disconnect_timeout {
messages.push(PeerMessage::Disconnected(conn.clone(), None)); messages.push(PeerMessage::Disconnected(conn.clone(), None, reason));
false false
} else { } else {
true true
} }
} }
ConnectionState::Error => { ConnectionState::Error => {
messages.push(PeerMessage::Disconnected(conn.clone(), conn.error.take())); messages.push(PeerMessage::Disconnected(
conn.clone(),
conn.error.take(),
CloseReason::Error,
));
false false
} }
_ => true, _ => true,
@ -584,9 +597,9 @@ pub enum ConnectionState {
Connected, Connected,
/// Connection is being closed /// Connection is being closed
Closing, Closing(CloseReason),
/// Connection is being closed by the other peer /// Connection is being closed by the other peer
ReceivingClosing, ReceivingClosing(CloseReason),
/// Connection has ended up in an erronous state /// Connection has ended up in an erronous state
Error, Error,

3
src/lib.rs
View File

@ -16,6 +16,7 @@ use thiserror::*;
use crate::{ use crate::{
connections::{Connection, ConnectionError, ConnectionManager}, connections::{Connection, ConnectionError, ConnectionManager},
listener::{Listener, ListenerMessage}, listener::{Listener, ListenerMessage},
package::CloseReason,
stats::NetStats, stats::NetStats,
}; };
@ -38,7 +39,7 @@ pub enum PeerMessage<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'st
/// A new connection has been connected /// A new connection has been connected
NewConnection(Connection<T>), NewConnection(Connection<T>),
/// An existing connection has disconnected, with an optional error /// An existing connection has disconnected, with an optional error
Disconnected(Connection<T>, Option<ConnectionError>), Disconnected(Connection<T>, Option<ConnectionError>, CloseReason),
/// The Peer has been closed /// The Peer has been closed
Closed, Closed,
/// A single message of type T /// A single message of type T

12
src/package.rs
View File

@ -5,10 +5,20 @@ pub enum Package<T> {
Hello(String), Hello(String),
Ping, Ping,
Pong, Pong,
Close, Close(CloseReason),
Messages(Messages<T>), Messages(Messages<T>),
} }
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
pub enum CloseReason {
NotAcceptingConnections,
Kicked,
Error,
ShuttingDown,
/// Closed because local peer initiated close
SelfClosed,
}
#[derive(Debug, Serialize, Deserialize)] #[derive(Debug, Serialize, Deserialize)]
pub struct Messages<T> { pub struct Messages<T> {
pub ack: u64, pub ack: u64,