add bandwidth tracking

This commit is contained in:
John Smith
2022-03-20 10:52:03 -04:00
parent 3888a832a0
commit ac0280e0b6
13 changed files with 203 additions and 107 deletions

View File

@@ -261,7 +261,8 @@ impl Network {
dial_info: DialInfo,
data: Vec<u8>,
) -> Result<(), String> {
match dial_info.protocol_type() {
let data_len = data.len();
let res = match dial_info.protocol_type() {
ProtocolType::UDP => {
let peer_socket_addr = dial_info.to_socket_addr();
RawUdpProtocolHandler::send_unbound_message(peer_socket_addr, data)
@@ -275,11 +276,17 @@ impl Network {
.map_err(logthru_net!())
}
ProtocolType::WS | ProtocolType::WSS => {
WebsocketProtocolHandler::send_unbound_message(dial_info, data)
WebsocketProtocolHandler::send_unbound_message(dial_info.clone(), data)
.await
.map_err(logthru_net!())
}
};
if res.is_ok() {
// Network accounting
self.network_manager()
.stats_packet_sent(dial_info.to_ip_addr(), data_len as u64);
}
res
}
async fn send_data_to_existing_connection(
@@ -287,6 +294,8 @@ impl Network {
descriptor: ConnectionDescriptor,
data: Vec<u8>,
) -> Result<Option<Vec<u8>>, String> {
let data_len = data.len();
// Handle connectionless protocol
if descriptor.protocol_type() == ProtocolType::UDP {
// send over the best udp socket we have bound since UDP is not connection oriented
@@ -299,6 +308,11 @@ impl Network {
.send_message(data, peer_socket_addr)
.await
.map_err(logthru_net!())?;
// Network accounting
self.network_manager()
.stats_packet_sent(peer_socket_addr.ip(), data_len as u64);
// Data was consumed
return Ok(None);
}
@@ -311,6 +325,10 @@ impl Network {
// connection exists, send over it
conn.send(data).await.map_err(logthru_net!())?;
// Network accounting
self.network_manager()
.stats_packet_sent(descriptor.remote.to_socket_addr().ip(), data_len as u64);
// Data was consumed
Ok(None)
} else {
@@ -326,14 +344,21 @@ impl Network {
dial_info: DialInfo,
data: Vec<u8>,
) -> Result<(), String> {
let data_len = data.len();
// Handle connectionless protocol
if dial_info.protocol_type() == ProtocolType::UDP {
let peer_socket_addr = dial_info.to_socket_addr();
if let Some(ph) = self.find_best_udp_protocol_handler(&peer_socket_addr, &None) {
return ph
let res = ph
.send_message(data, peer_socket_addr)
.await
.map_err(logthru_net!());
if res.is_ok() {
// Network accounting
self.network_manager()
.stats_packet_sent(peer_socket_addr.ip(), data_len as u64);
}
return res;
}
return Err("no appropriate UDP protocol handler for dial_info".to_owned())
.map_err(logthru_net!(error));
@@ -343,10 +368,16 @@ impl Network {
let local_addr = self.get_preferred_local_address(&dial_info);
let conn = self
.connection_manager()
.get_or_create_connection(Some(local_addr), dial_info)
.get_or_create_connection(Some(local_addr), dial_info.clone())
.await?;
conn.send(data).await.map_err(logthru_net!(error))
let res = conn.send(data).await.map_err(logthru_net!(error));
if res.is_ok() {
// Network accounting
self.network_manager()
.stats_packet_sent(dial_info.to_ip_addr(), data_len as u64);
}
res
}
// Send data to node

View File

@@ -54,6 +54,13 @@ impl Network {
// XXX: Limit the number of packets from the same IP address?
log_net!("UDP packet: {:?}", descriptor);
// Network accounting
network_manager.stats_packet_rcvd(
descriptor.remote.to_socket_addr().ip(),
size as u64,
);
// Pass it up for processing
if let Err(e) = network_manager
.on_recv_envelope(&data[..size], descriptor)
.await

View File

@@ -49,6 +49,32 @@ impl Network {
/////////////////////////////////////////////////////////////////
pub async fn send_data_unbound_to_dial_info(
&self,
dial_info: DialInfo,
data: Vec<u8>,
) -> Result<(), String> {
let res = match dial_info.protocol_type() {
ProtocolType::UDP => {
return Err("no support for UDP protocol".to_owned()).map_err(logthru_net!(error))
}
ProtocolType::TCP => {
return Err("no support for TCP protocol".to_owned()).map_err(logthru_net!(error))
}
ProtocolType::WS | ProtocolType::WSS => {
WebsocketProtocolHandler::send_unbound_message(dial_info, data)
.await
.map_err(logthru_net!())
}
};
if res.is_ok() {
// Network accounting
self.network_manager()
.stats_packet_sent(dial_info.to_ip_addr(), data_len as u64);
}
res
}
async fn send_data_to_existing_connection(
&self,
descriptor: ConnectionDescriptor,
@@ -71,6 +97,10 @@ impl Network {
// connection exists, send over it
conn.send(data).await.map_err(logthru_net!())?;
// Network accounting
self.network_manager()
.stats_packet_sent(descriptor.remote.to_socket_addr().ip(), data_len as u64);
// Data was consumed
Ok(None)
} else {
@@ -80,26 +110,6 @@ impl Network {
}
}
pub async fn send_data_unbound_to_dial_info(
&self,
dial_info: DialInfo,
data: Vec<u8>,
) -> Result<(), String> {
match dial_info.protocol_type() {
ProtocolType::UDP => {
return Err("no support for UDP protocol".to_owned()).map_err(logthru_net!(error))
}
ProtocolType::TCP => {
return Err("no support for TCP protocol".to_owned()).map_err(logthru_net!(error))
}
ProtocolType::WS | ProtocolType::WSS => {
WebsocketProtocolHandler::send_unbound_message(dial_info, data)
.await
.map_err(logthru_net!())
}
}
}
pub async fn send_data_to_dial_info(
&self,
dial_info: DialInfo,
@@ -118,7 +128,13 @@ impl Network {
.get_or_create_connection(None, dial_info)
.await?;
conn.send(data).await.map_err(logthru_net!(error))
let res = conn.send(data).await.map_err(logthru_net!(error));
if res.is_ok() {
// Network accounting
self.network_manager()
.stats_packet_sent(dial_info.to_ip_addr(), data_len as u64);
}
res
}
pub async fn send_data(&self, node_ref: NodeRef, data: Vec<u8>) -> Result<(), String> {