Compare commits

...
Author SHA1 Message Date
fanyang 82c66acc23 fix: bound peer rpc packet queues 2026-06-23 21:33:37 +08:00
fanyang 65f487ba26 fix: make stats counters thread safe 2026-06-23 21:18:07 +08:00
4 changed files with 342 additions and 151 deletions
+75 -103
View File
@@ -1,8 +1,9 @@
use dashmap::DashMap;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use std::cell::UnsafeCell;
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::time::interval;
use tokio_util::task::AbortOnDropHandle;
@@ -22,6 +23,8 @@ pub enum MetricName {
PeerRpcDuration,
/// RPC errors
PeerRpcErrors,
/// RPC/control packets dropped because the peer RPC queue is unavailable
PeerRpcPacketQueueDrops,
/// Data-plane traffic bytes sent
TrafficBytesTx,
@@ -115,6 +118,7 @@ impl fmt::Display for MetricName {
MetricName::PeerRpcServerRx => write!(f, "peer_rpc_server_rx"),
MetricName::PeerRpcDuration => write!(f, "peer_rpc_duration_ms"),
MetricName::PeerRpcErrors => write!(f, "peer_rpc_errors"),
MetricName::PeerRpcPacketQueueDrops => write!(f, "peer_rpc_packet_queue_drops"),
MetricName::TrafficBytesTx => write!(f, "traffic_bytes_tx"),
MetricName::TrafficBytesTxByInstance => write!(f, "traffic_bytes_tx_by_instance"),
@@ -374,10 +378,10 @@ impl Default for LabelSet {
}
}
/// UnsafeCounter provides a high-performance counter using UnsafeCell
/// UnsafeCounter provides a high-performance atomic counter
#[derive(Debug)]
pub struct UnsafeCounter {
value: UnsafeCell<u64>,
value: AtomicU64,
}
impl Default for UnsafeCounter {
@@ -389,121 +393,79 @@ impl Default for UnsafeCounter {
impl UnsafeCounter {
pub fn new() -> Self {
Self {
value: UnsafeCell::new(0),
value: AtomicU64::new(0),
}
}
pub fn new_with_value(initial: u64) -> Self {
Self {
value: UnsafeCell::new(initial),
value: AtomicU64::new(initial),
}
}
/// Increment the counter by the given amount
/// # Safety
/// This method is unsafe because it uses UnsafeCell. The caller must ensure
/// that no other thread is accessing this counter simultaneously.
pub unsafe fn add(&self, delta: u64) {
let ptr = self.value.get();
unsafe {
*ptr = (*ptr).saturating_add(delta);
}
pub fn add(&self, delta: u64) {
let _ = self
.value
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
Some(current.saturating_add(delta))
});
}
/// Increment the counter by 1
/// # Safety
/// This method is unsafe because it uses UnsafeCell. The caller must ensure
/// that no other thread is accessing this counter simultaneously.
pub unsafe fn inc(&self) {
unsafe {
self.add(1);
}
pub fn inc(&self) {
self.add(1);
}
/// Get the current value of the counter
/// # Safety
/// This method is unsafe because it uses UnsafeCell. The caller must ensure
/// that no other thread is modifying this counter simultaneously.
pub unsafe fn get(&self) -> u64 {
let ptr = self.value.get();
unsafe { *ptr }
pub fn get(&self) -> u64 {
self.value.load(Ordering::Relaxed)
}
/// Reset the counter to zero
/// # Safety
/// This method is unsafe because it uses UnsafeCell. The caller must ensure
/// that no other thread is accessing this counter simultaneously.
pub unsafe fn reset(&self) {
let ptr = self.value.get();
unsafe {
*ptr = 0;
}
pub fn reset(&self) {
self.value.store(0, Ordering::Relaxed);
}
/// Set the counter to a specific value
/// # Safety
/// This method is unsafe because it uses UnsafeCell. The caller must ensure
/// that no other thread is accessing this counter simultaneously.
pub unsafe fn set(&self, value: u64) {
let ptr = self.value.get();
unsafe {
*ptr = value;
}
pub fn set(&self, value: u64) {
self.value.store(value, Ordering::Relaxed);
}
}
// UnsafeCounter is Send + Sync because the safety is guaranteed by the caller
unsafe impl Send for UnsafeCounter {}
unsafe impl Sync for UnsafeCounter {}
/// MetricData contains both the counter and last update timestamp
/// Uses UnsafeCell for lock-free access
#[derive(Debug)]
struct MetricData {
counter: UnsafeCounter,
last_updated: UnsafeCell<Instant>,
last_updated: Mutex<Instant>,
}
impl MetricData {
fn new() -> Self {
Self {
counter: UnsafeCounter::new(),
last_updated: UnsafeCell::new(Instant::now()),
last_updated: Mutex::new(Instant::now()),
}
}
fn new_with_value(initial: u64) -> Self {
Self {
counter: UnsafeCounter::new_with_value(initial),
last_updated: UnsafeCell::new(Instant::now()),
last_updated: Mutex::new(Instant::now()),
}
}
/// Update the last_updated timestamp
/// # Safety
/// This method is unsafe because it uses UnsafeCell. The caller must ensure
/// that no other thread is accessing this timestamp simultaneously.
unsafe fn touch(&self) {
let ptr = self.last_updated.get();
unsafe {
*ptr = Instant::now();
}
fn touch(&self) {
*self.last_updated.lock() = Instant::now();
}
/// Get the last updated timestamp
/// # Safety
/// This method is unsafe because it uses UnsafeCell. The caller must ensure
/// that no other thread is modifying this timestamp simultaneously.
unsafe fn get_last_updated(&self) -> Instant {
let ptr = self.last_updated.get();
unsafe { *ptr }
fn get_last_updated(&self) -> Instant {
*self.last_updated.lock()
}
}
// MetricData is Send + Sync because the safety is guaranteed by the caller
unsafe impl Send for MetricData {}
unsafe impl Sync for MetricData {}
/// MetricKey uniquely identifies a metric with its name and labels
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct MetricKey {
@@ -546,39 +508,31 @@ impl CounterHandle {
/// Increment the counter by the given amount
pub fn add(&self, delta: u64) {
unsafe {
self.metric_data.counter.add(delta);
self.metric_data.touch();
}
self.metric_data.counter.add(delta);
self.metric_data.touch();
}
/// Increment the counter by 1
pub fn inc(&self) {
unsafe {
self.metric_data.counter.inc();
self.metric_data.touch();
}
self.metric_data.counter.inc();
self.metric_data.touch();
}
/// Get the current value of the counter
pub fn get(&self) -> u64 {
unsafe { self.metric_data.counter.get() }
self.metric_data.counter.get()
}
/// Reset the counter to zero
pub fn reset(&self) {
unsafe {
self.metric_data.counter.reset();
self.metric_data.touch();
}
self.metric_data.counter.reset();
self.metric_data.touch();
}
/// Set the counter to a specific value
pub fn set(&self, value: u64) {
unsafe {
self.metric_data.counter.set(value);
self.metric_data.touch();
}
self.metric_data.counter.set(value);
self.metric_data.touch();
}
}
@@ -624,7 +578,7 @@ impl StatsManager {
counters.retain(|_, metric_data: &mut Arc<MetricData>| {
Arc::strong_count(metric_data) > 1
|| unsafe { metric_data.get_last_updated() > cutoff_time }
|| metric_data.get_last_updated() > cutoff_time
});
counters.shrink_to_fit();
}
@@ -662,7 +616,7 @@ impl StatsManager {
let key = entry.key();
let metric_data = entry.value();
let value = unsafe { metric_data.counter.get() };
let value = metric_data.counter.get();
metrics.push(MetricSnapshot {
name: key.name,
@@ -695,7 +649,7 @@ impl StatsManager {
let key = MetricKey::new(name, labels.clone());
if let Some(metric_data) = self.counters.get(&key) {
let value = unsafe { metric_data.counter.get() };
let value = metric_data.counter.get();
Some(MetricSnapshot {
name,
labels: labels.clone(),
@@ -796,17 +750,15 @@ mod tests {
async fn test_unsafe_counter() {
let counter = UnsafeCounter::new();
unsafe {
assert_eq!(counter.get(), 0);
counter.inc();
assert_eq!(counter.get(), 1);
counter.add(5);
assert_eq!(counter.get(), 6);
counter.set(10);
assert_eq!(counter.get(), 10);
counter.reset();
assert_eq!(counter.get(), 0);
}
assert_eq!(counter.get(), 0);
counter.inc();
assert_eq!(counter.get(), 1);
counter.add(5);
assert_eq!(counter.get(), 6);
counter.set(10);
assert_eq!(counter.get(), 10);
counter.reset();
assert_eq!(counter.get(), 0);
}
#[tokio::test]
@@ -951,8 +903,7 @@ mod tests {
stats
.counters
.retain(|_, metric_data: &mut Arc<MetricData>| {
Arc::strong_count(metric_data) > 1
|| unsafe { metric_data.get_last_updated() > cutoff_time }
Arc::strong_count(metric_data) > 1 || metric_data.get_last_updated() > cutoff_time
});
assert_eq!(stats.metric_count(), 1);
@@ -962,12 +913,33 @@ mod tests {
stats
.counters
.retain(|_, metric_data: &mut Arc<MetricData>| {
Arc::strong_count(metric_data) > 1
|| unsafe { metric_data.get_last_updated() > cutoff_time }
Arc::strong_count(metric_data) > 1 || metric_data.get_last_updated() > cutoff_time
});
assert_eq!(stats.metric_count(), 0);
}
#[tokio::test]
async fn test_counter_handle_concurrent_increment() {
const THREADS: usize = 8;
const INCREMENTS_PER_THREAD: usize = 10_000;
let stats = StatsManager::new();
let counter = stats.get_simple_counter(MetricName::TrafficPacketsForwarded);
std::thread::scope(|scope| {
for _ in 0..THREADS {
let counter = counter.clone();
scope.spawn(move || {
for _ in 0..INCREMENTS_PER_THREAD {
counter.inc();
}
});
}
});
assert_eq!(counter.get(), (THREADS * INCREMENTS_PER_THREAD) as u64);
}
#[tokio::test]
async fn test_stats_rpc_data_structures() {
// Test GetStatsRequest
+100 -8
View File
@@ -18,7 +18,7 @@ use guarden::{Guard, defer};
use tokio::{
sync::{
Mutex,
mpsc::{self, UnboundedReceiver, UnboundedSender},
mpsc::{self, Receiver, Sender, error::TrySendError},
},
task::JoinSet,
};
@@ -30,7 +30,7 @@ use crate::{
error::Error,
global_ctx::{ArcGlobalCtx, GlobalCtx, GlobalCtxEvent, NetworkIdentity, TrustedKeySource},
join_joinset_background, shrink_dashmap,
stats_manager::{LabelSet, LabelType, MetricName, StatsManager},
stats_manager::{CounterHandle, LabelSet, LabelType, MetricName, StatsManager},
token_bucket::TokenBucket,
},
peer_center::instance::{PeerCenterInstance, PeerMapWithPeerRpcManager},
@@ -64,6 +64,35 @@ use super::{
},
};
const PEER_RPC_PACKET_QUEUE_CAPACITY: usize = 1024;
fn try_enqueue_peer_rpc_packet(
sender: &Sender<ZCPacket>,
packet: ZCPacket,
dropped_packets: &CounterHandle,
queue_name: &'static str,
) -> bool {
match sender.try_send(packet) {
Ok(()) => true,
Err(TrySendError::Full(_)) => {
dropped_packets.inc();
tracing::warn!(
queue = queue_name,
"drop peer rpc/control packet because queue is full"
);
false
}
Err(TrySendError::Closed(_)) => {
dropped_packets.inc();
tracing::warn!(
queue = queue_name,
"drop peer rpc/control packet because receiver is closed"
);
false
}
}
}
#[async_trait::async_trait]
#[auto_impl::auto_impl(&, Box, Arc)]
pub trait GlobalForeignNetworkAccessor: Send + Sync + 'static {
@@ -87,7 +116,7 @@ struct ForeignNetworkEntry {
pm_packet_sender: Mutex<Option<PacketRecvChan>>,
peer_rpc: Arc<PeerRpcManager>,
rpc_sender: UnboundedSender<ZCPacket>,
rpc_sender: Sender<ZCPacket>,
packet_recv: Mutex<Option<PacketRecvChanReceiver>>,
@@ -312,12 +341,12 @@ impl ForeignNetworkEntry {
fn build_rpc_tspt(
my_peer_id: PeerId,
peer_map: Arc<PeerMap>,
) -> (Arc<PeerRpcManager>, UnboundedSender<ZCPacket>) {
) -> (Arc<PeerRpcManager>, Sender<ZCPacket>) {
struct RpcTransport {
my_peer_id: PeerId,
peer_map: Weak<PeerMap>,
packet_recv: Mutex<UnboundedReceiver<ZCPacket>>,
packet_recv: Mutex<Receiver<ZCPacket>>,
}
#[async_trait::async_trait]
@@ -359,7 +388,8 @@ impl ForeignNetworkEntry {
}
}
let (rpc_transport_sender, peer_rpc_tspt_recv) = mpsc::unbounded_channel();
let (rpc_transport_sender, peer_rpc_tspt_recv) =
mpsc::channel(PEER_RPC_PACKET_QUEUE_CAPACITY);
let tspt = RpcTransport {
my_peer_id,
peer_map: Arc::downgrade(&peer_map),
@@ -478,6 +508,9 @@ impl ForeignNetworkEntry {
let rx_packets = self
.stats_mgr
.get_counter(MetricName::TrafficPacketsRx, label_set.clone());
let rpc_queue_drops = self
.stats_mgr
.get_counter(MetricName::PeerRpcPacketQueueDrops, label_set.clone());
self.tasks.lock().await.spawn(async move {
while let Ok(mut zc_packet) = recv_packet_from_chan(&mut recv).await {
@@ -526,7 +559,12 @@ impl ForeignNetworkEntry {
{
rx_bytes.add(buf_len as u64);
rx_packets.inc();
rpc_sender.send(zc_packet).unwrap();
try_enqueue_peer_rpc_packet(
&rpc_sender,
zc_packet,
&rpc_queue_drops,
"foreign_network_peer_rpc",
);
continue;
}
tracing::trace!(
@@ -1236,7 +1274,7 @@ impl Drop for ForeignNetworkManager {
pub mod tests {
use crate::{
common::global_ctx::tests::get_mock_global_ctx_with_network,
common::stats_manager::{LabelSet, LabelType, MetricName},
common::stats_manager::{LabelSet, LabelType, MetricName, StatsManager},
connector::udp_hole_punch::tests::{
create_mock_peer_manager_with_mock_stun, replace_stun_info_collector,
},
@@ -1253,6 +1291,7 @@ pub mod tests {
},
};
use std::{collections::HashMap, time::Duration};
use tokio::sync::mpsc;
use super::*;
@@ -1265,6 +1304,59 @@ pub mod tests {
.unwrap_or(0)
}
#[tokio::test]
async fn peer_rpc_queue_helper_enqueues_when_available() {
let (sender, mut receiver) = mpsc::channel(1);
let stats_manager = StatsManager::new();
let dropped_packets = stats_manager.get_simple_counter(MetricName::PeerRpcPacketQueueDrops);
assert!(try_enqueue_peer_rpc_packet(
&sender,
ZCPacket::new_with_payload(b"rpc"),
&dropped_packets,
"test_foreign_peer_rpc",
));
assert_eq!(dropped_packets.get(), 0);
assert!(receiver.try_recv().is_ok());
}
#[tokio::test]
async fn peer_rpc_queue_helper_drops_when_full() {
let (sender, _receiver) = mpsc::channel(1);
sender
.try_send(ZCPacket::new_with_payload(b"existing"))
.unwrap();
let stats_manager = StatsManager::new();
let dropped_packets = stats_manager.get_simple_counter(MetricName::PeerRpcPacketQueueDrops);
assert!(!try_enqueue_peer_rpc_packet(
&sender,
ZCPacket::new_with_payload(b"overflow"),
&dropped_packets,
"test_foreign_peer_rpc",
));
assert_eq!(dropped_packets.get(), 1);
}
#[tokio::test]
async fn peer_rpc_queue_helper_drops_when_closed() {
let (sender, receiver) = mpsc::channel(1);
drop(receiver);
let stats_manager = StatsManager::new();
let dropped_packets = stats_manager.get_simple_counter(MetricName::PeerRpcPacketQueueDrops);
assert!(!try_enqueue_peer_rpc_packet(
&sender,
ZCPacket::new_with_payload(b"closed"),
&dropped_packets,
"test_foreign_peer_rpc",
));
assert_eq!(dropped_packets.get(), 1);
}
async fn create_mock_peer_manager_for_foreign_network_ext(
network: &str,
secret: &str,
+106 -8
View File
@@ -13,7 +13,7 @@ use std::{
use tokio::{
sync::{
Mutex, RwLock,
mpsc::{self, UnboundedReceiver, UnboundedSender},
mpsc::{self, Receiver, Sender, error::TrySendError},
},
task::JoinSet,
};
@@ -72,14 +72,43 @@ use super::{
route_trait::{ArcRoute, Route},
};
const PEER_RPC_PACKET_QUEUE_CAPACITY: usize = 1024;
fn try_enqueue_peer_rpc_packet(
sender: &Sender<ZCPacket>,
packet: ZCPacket,
dropped_packets: &CounterHandle,
queue_name: &'static str,
) -> bool {
match sender.try_send(packet) {
Ok(()) => true,
Err(TrySendError::Full(_)) => {
dropped_packets.inc();
tracing::warn!(
queue = queue_name,
"drop peer rpc/control packet because queue is full"
);
false
}
Err(TrySendError::Closed(_)) => {
dropped_packets.inc();
tracing::warn!(
queue = queue_name,
"drop peer rpc/control packet because receiver is closed"
);
false
}
}
}
struct RpcTransport {
my_peer_id: PeerId,
peers: Weak<PeerMap>,
// TODO: this seems can be removed
foreign_peers: Mutex<Option<Weak<ForeignNetworkClient>>>,
packet_recv: Mutex<UnboundedReceiver<ZCPacket>>,
peer_rpc_tspt_sender: UnboundedSender<ZCPacket>,
packet_recv: Mutex<Receiver<ZCPacket>>,
peer_rpc_tspt_sender: Sender<ZCPacket>,
encryptor: Arc<dyn Encryptor>,
is_secure_mode_enabled: bool,
@@ -273,7 +302,8 @@ impl PeerManager {
.unwrap_or(false);
// TODO: remove these because we have impl pipeline processor.
let (peer_rpc_tspt_sender, peer_rpc_tspt_recv) = mpsc::unbounded_channel();
let (peer_rpc_tspt_sender, peer_rpc_tspt_recv) =
mpsc::channel(PEER_RPC_PACKET_QUEUE_CAPACITY);
let rpc_tspt = Arc::new(RpcTransport {
my_peer_id,
peers: Arc::downgrade(&peers),
@@ -1243,7 +1273,8 @@ impl PeerManager {
// for peer rpc packet
struct PeerRpcPacketProcessor {
peer_rpc_tspt_sender: UnboundedSender<ZCPacket>,
peer_rpc_tspt_sender: Sender<ZCPacket>,
dropped_packets: CounterHandle,
}
#[async_trait::async_trait]
@@ -1254,15 +1285,27 @@ impl PeerManager {
|| hdr.packet_type == PacketType::RpcReq as u8
|| hdr.packet_type == PacketType::RpcResp as u8
{
self.peer_rpc_tspt_sender.send(packet).unwrap();
try_enqueue_peer_rpc_packet(
&self.peer_rpc_tspt_sender,
packet,
&self.dropped_packets,
"local_peer_rpc",
);
None
} else {
Some(packet)
}
}
}
let peer_rpc_queue_drops = self.global_ctx.stats_manager().get_counter(
MetricName::PeerRpcPacketQueueDrops,
LabelSet::new().with_label_type(LabelType::NetworkName(
self.global_ctx.get_network_name().to_string(),
)),
);
self.add_packet_process_pipeline(Box::new(PeerRpcPacketProcessor {
peer_rpc_tspt_sender: self.peer_rpc_tspt.peer_rpc_tspt_sender.clone(),
dropped_packets: peer_rpc_queue_drops,
}))
.await;
}
@@ -2209,7 +2252,7 @@ mod tests {
PeerId,
config::Flags,
global_ctx::{NetworkIdentity, tests::get_mock_global_ctx},
stats_manager::{LabelSet, LabelType, MetricName},
stats_manager::{LabelSet, LabelType, MetricName, StatsManager},
},
connector::{
create_connector_by_url, direct::PeerManagerForDirectConnector,
@@ -2240,7 +2283,9 @@ mod tests {
},
};
use super::PeerManager;
use tokio::sync::mpsc;
use super::{PeerManager, try_enqueue_peer_rpc_packet};
async fn create_lazy_peer_manager() -> Arc<PeerManager> {
let peer_mgr = create_mock_peer_manager_with_mock_stun(NatType::Unknown).await;
@@ -2265,6 +2310,59 @@ mod tests {
))
}
#[tokio::test]
async fn peer_rpc_queue_helper_enqueues_when_available() {
let (sender, mut receiver) = mpsc::channel(1);
let stats_manager = StatsManager::new();
let dropped_packets = stats_manager.get_simple_counter(MetricName::PeerRpcPacketQueueDrops);
assert!(try_enqueue_peer_rpc_packet(
&sender,
ZCPacket::new_with_payload(b"rpc"),
&dropped_packets,
"test_peer_rpc",
));
assert_eq!(dropped_packets.get(), 0);
assert!(receiver.try_recv().is_ok());
}
#[tokio::test]
async fn peer_rpc_queue_helper_drops_when_full() {
let (sender, _receiver) = mpsc::channel(1);
sender
.try_send(ZCPacket::new_with_payload(b"existing"))
.unwrap();
let stats_manager = StatsManager::new();
let dropped_packets = stats_manager.get_simple_counter(MetricName::PeerRpcPacketQueueDrops);
assert!(!try_enqueue_peer_rpc_packet(
&sender,
ZCPacket::new_with_payload(b"overflow"),
&dropped_packets,
"test_peer_rpc",
));
assert_eq!(dropped_packets.get(), 1);
}
#[tokio::test]
async fn peer_rpc_queue_helper_drops_when_closed() {
let (sender, receiver) = mpsc::channel(1);
drop(receiver);
let stats_manager = StatsManager::new();
let dropped_packets = stats_manager.get_simple_counter(MetricName::PeerRpcPacketQueueDrops);
assert!(!try_enqueue_peer_rpc_packet(
&sender,
ZCPacket::new_with_payload(b"closed"),
&dropped_packets,
"test_peer_rpc",
));
assert_eq!(dropped_packets.get(), 1);
}
struct TestCostCalculator {
costs: HashMap<(PeerId, PeerId), i32>,
}
+61 -32
View File
@@ -1,7 +1,4 @@
use std::{
cell::UnsafeCell,
sync::atomic::{AtomicU32, Ordering::Relaxed},
};
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering::Relaxed};
pub struct WindowLatency {
latency_us_window: Vec<AtomicU32>,
@@ -63,34 +60,30 @@ impl WindowLatency {
#[derive(Debug)]
pub struct Throughput {
tx_bytes: UnsafeCell<u64>,
rx_bytes: UnsafeCell<u64>,
tx_packets: UnsafeCell<u64>,
rx_packets: UnsafeCell<u64>,
tx_bytes: AtomicU64,
rx_bytes: AtomicU64,
tx_packets: AtomicU64,
rx_packets: AtomicU64,
}
impl Clone for Throughput {
fn clone(&self) -> Self {
Self {
tx_bytes: UnsafeCell::new(unsafe { *self.tx_bytes.get() }),
rx_bytes: UnsafeCell::new(unsafe { *self.rx_bytes.get() }),
tx_packets: UnsafeCell::new(unsafe { *self.tx_packets.get() }),
rx_packets: UnsafeCell::new(unsafe { *self.rx_packets.get() }),
tx_bytes: AtomicU64::new(self.tx_bytes()),
rx_bytes: AtomicU64::new(self.rx_bytes()),
tx_packets: AtomicU64::new(self.tx_packets()),
rx_packets: AtomicU64::new(self.rx_packets()),
}
}
}
// add sync::Send and sync::Sync traits to Throughput
unsafe impl Send for Throughput {}
unsafe impl Sync for Throughput {}
impl Default for Throughput {
fn default() -> Self {
Self {
tx_bytes: UnsafeCell::new(0),
rx_bytes: UnsafeCell::new(0),
tx_packets: UnsafeCell::new(0),
rx_packets: UnsafeCell::new(0),
tx_bytes: AtomicU64::new(0),
rx_bytes: AtomicU64::new(0),
tx_packets: AtomicU64::new(0),
rx_packets: AtomicU64::new(0),
}
}
}
@@ -101,32 +94,68 @@ impl Throughput {
}
pub fn tx_bytes(&self) -> u64 {
unsafe { *self.tx_bytes.get() }
self.tx_bytes.load(Relaxed)
}
pub fn rx_bytes(&self) -> u64 {
unsafe { *self.rx_bytes.get() }
self.rx_bytes.load(Relaxed)
}
pub fn tx_packets(&self) -> u64 {
unsafe { *self.tx_packets.get() }
self.tx_packets.load(Relaxed)
}
pub fn rx_packets(&self) -> u64 {
unsafe { *self.rx_packets.get() }
self.rx_packets.load(Relaxed)
}
pub fn record_tx_bytes(&self, bytes: u64) {
unsafe {
*self.tx_bytes.get() += bytes;
*self.tx_packets.get() += 1;
}
self.tx_bytes.fetch_add(bytes, Relaxed);
self.tx_packets.fetch_add(1, Relaxed);
}
pub fn record_rx_bytes(&self, bytes: u64) {
unsafe {
*self.rx_bytes.get() += bytes;
*self.rx_packets.get() += 1;
}
self.rx_bytes.fetch_add(bytes, Relaxed);
self.rx_packets.fetch_add(1, Relaxed);
}
}
#[cfg(test)]
mod tests {
use super::Throughput;
use std::sync::Arc;
#[test]
fn throughput_records_concurrent_tx_and_rx() {
const THREADS: usize = 8;
const RECORDS_PER_THREAD: usize = 10_000;
const TX_BYTES_PER_RECORD: u64 = 3;
const RX_BYTES_PER_RECORD: u64 = 7;
let throughput = Arc::new(Throughput::new());
std::thread::scope(|scope| {
for _ in 0..THREADS {
let throughput = Arc::clone(&throughput);
scope.spawn(move || {
for _ in 0..RECORDS_PER_THREAD {
throughput.record_tx_bytes(TX_BYTES_PER_RECORD);
throughput.record_rx_bytes(RX_BYTES_PER_RECORD);
}
});
}
});
let expected_packets = (THREADS * RECORDS_PER_THREAD) as u64;
assert_eq!(throughput.tx_packets(), expected_packets);
assert_eq!(throughput.rx_packets(), expected_packets);
assert_eq!(
throughput.tx_bytes(),
expected_packets * TX_BYTES_PER_RECORD
);
assert_eq!(
throughput.rx_bytes(),
expected_packets * RX_BYTES_PER_RECORD
);
}
}