Add close reasons

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

View File

@ -12,7 +12,7 @@ use thiserror::Error;
use crate::{
PeerConfig, PeerMessage,
listener::{DATAGRAM_SIZE, ListenerError},
package::{Message, Messages, Package},
package::{CloseReason, Message, Messages, Package},
stats::NetStats,
};
@ -71,7 +71,7 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
pub fn close(&mut self) {
self.closing_since = Some(Instant::now());
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 => {
match self.udp.send_to(*addr, Package::<T>::Close) {
ConnectionState::Closing(reason)
| ConnectionState::ReceivingClosing(reason) => {
match self.udp.send_to(*addr, Package::<T>::Close(reason)) {
Ok(bytes) => {
conn.bytes_tx += bytes;
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
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 conn.state != ConnectionState::Closing {
if !matches!(conn.state, ConnectionState::Closing(_)) {
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 {
if !self.accepting_connections || self.config.identifier != identifier {
self.connections
.insert(*addr, Connection::from(*addr, ConnectionState::Closing));
self.connections.insert(
*addr,
Connection::from(
*addr,
ConnectionState::Closing(CloseReason::NotAcceptingConnections),
),
);
} else {
self.connections.insert(
*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) {
match conn.state {
ConnectionState::Closing => {
messages
.push(PeerMessage::Disconnected(conn.clone(), conn.error.take()));
ConnectionState::Closing(_) => {
messages.push(PeerMessage::Disconnected(
conn.clone(),
conn.error.take(),
CloseReason::SelfClosed,
));
true
}
ConnectionState::Error => false,
ConnectionState::ReceivingClosing => {
ConnectionState::ReceivingClosing(_) => {
conn.last_recv_close = Instant::now();
false
}
_ => {
conn.last_recv_close = Instant::now();
conn.state = ConnectionState::ReceivingClosing;
conn.state = ConnectionState::ReceivingClosing(reason);
false
}
}
@ -401,16 +410,20 @@ impl<T: std::fmt::Debug + Clone + Serialize + DeserializeOwned + Send + Sync + '
self.connections.retain(|_, conn| {
let retain = match conn.state {
ConnectionState::ReceivingClosing => {
ConnectionState::ReceivingClosing(reason) => {
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
} else {
true
}
}
ConnectionState::Error => {
messages.push(PeerMessage::Disconnected(conn.clone(), conn.error.take()));
messages.push(PeerMessage::Disconnected(
conn.clone(),
conn.error.take(),
CloseReason::Error,
));
false
}
_ => true,
@ -584,9 +597,9 @@ pub enum ConnectionState {
Connected,
/// Connection is being closed
Closing,
Closing(CloseReason),
/// Connection is being closed by the other peer
ReceivingClosing,
ReceivingClosing(CloseReason),
/// Connection has ended up in an erronous state
Error,

View File

@ -16,6 +16,7 @@ use thiserror::*;
use crate::{
connections::{Connection, ConnectionError, ConnectionManager},
listener::{Listener, ListenerMessage},
package::CloseReason,
stats::NetStats,
};
@ -38,7 +39,7 @@ pub enum PeerMessage<T: Clone + Serialize + DeserializeOwned + Send + Sync + 'st
/// A new connection has been connected
NewConnection(Connection<T>),
/// 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
Closed,
/// A single message of type T

View File

@ -5,10 +5,20 @@ pub enum Package<T> {
Hello(String),
Ping,
Pong,
Close,
Close(CloseReason),
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)]
pub struct Messages<T> {
pub ack: u64,