mirror of
https://github.com/EasyTier/EasyTier.git
synced 2026-09-21 03:24:00 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
82c66acc23 | ||
|
|
65f487ba26 |
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>,
|
||||
}
|
||||
|
||||
@@ -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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user