mirror of https://github.com/telemt/telemt.git
Compare commits
7 Commits
7b25f62438
...
14804d4f74
| Author | SHA1 | Date |
|---|---|---|
|
|
14804d4f74 | |
|
|
8ac1a0017d | |
|
|
3df274caa6 | |
|
|
780546a680 | |
|
|
729ffa0fcd | |
|
|
e594d6f079 | |
|
|
ecd6a19246 |
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "telemt"
|
||||
version = "3.3.35"
|
||||
version = "3.3.36"
|
||||
edition = "2024"
|
||||
|
||||
[features]
|
||||
|
|
|
|||
|
|
@ -26,6 +26,15 @@ pub struct UserIpTracker {
|
|||
cleanup_drain_lock: Arc<AsyncMutex<()>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct UserIpTrackerMemoryStats {
|
||||
pub active_users: usize,
|
||||
pub recent_users: usize,
|
||||
pub active_entries: usize,
|
||||
pub recent_entries: usize,
|
||||
pub cleanup_queue_len: usize,
|
||||
}
|
||||
|
||||
impl UserIpTracker {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
|
|
@ -141,6 +150,13 @@ impl UserIpTracker {
|
|||
|
||||
let mut active_ips = self.active_ips.write().await;
|
||||
let mut recent_ips = self.recent_ips.write().await;
|
||||
let window = *self.limit_window.read().await;
|
||||
let now = Instant::now();
|
||||
|
||||
for user_recent in recent_ips.values_mut() {
|
||||
Self::prune_recent(user_recent, now, window);
|
||||
}
|
||||
|
||||
let mut users =
|
||||
Vec::<String>::with_capacity(active_ips.len().saturating_add(recent_ips.len()));
|
||||
users.extend(active_ips.keys().cloned());
|
||||
|
|
@ -166,6 +182,26 @@ impl UserIpTracker {
|
|||
}
|
||||
}
|
||||
|
||||
pub async fn memory_stats(&self) -> UserIpTrackerMemoryStats {
|
||||
let cleanup_queue_len = self
|
||||
.cleanup_queue
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||
.len();
|
||||
let active_ips = self.active_ips.read().await;
|
||||
let recent_ips = self.recent_ips.read().await;
|
||||
let active_entries = active_ips.values().map(HashMap::len).sum();
|
||||
let recent_entries = recent_ips.values().map(HashMap::len).sum();
|
||||
|
||||
UserIpTrackerMemoryStats {
|
||||
active_users: active_ips.len(),
|
||||
recent_users: recent_ips.len(),
|
||||
active_entries,
|
||||
recent_entries,
|
||||
cleanup_queue_len,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn set_limit_policy(&self, mode: UserMaxUniqueIpsMode, window_secs: u64) {
|
||||
{
|
||||
let mut current_mode = self.limit_mode.write().await;
|
||||
|
|
@ -451,6 +487,7 @@ impl Default for UserIpTracker {
|
|||
mod tests {
|
||||
use super::*;
|
||||
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr {
|
||||
IpAddr::V4(Ipv4Addr::new(oct1, oct2, oct3, oct4))
|
||||
|
|
@ -764,4 +801,54 @@ mod tests {
|
|||
tokio::time::sleep(Duration::from_millis(1100)).await;
|
||||
assert!(tracker.check_and_add("test_user", ip2).await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_memory_stats_reports_queue_and_entry_counts() {
|
||||
let tracker = UserIpTracker::new();
|
||||
tracker.set_user_limit("test_user", 4).await;
|
||||
let ip1 = test_ipv4(10, 2, 0, 1);
|
||||
let ip2 = test_ipv4(10, 2, 0, 2);
|
||||
|
||||
tracker.check_and_add("test_user", ip1).await.unwrap();
|
||||
tracker.check_and_add("test_user", ip2).await.unwrap();
|
||||
tracker.enqueue_cleanup("test_user".to_string(), ip1);
|
||||
|
||||
let snapshot = tracker.memory_stats().await;
|
||||
assert_eq!(snapshot.active_users, 1);
|
||||
assert_eq!(snapshot.recent_users, 1);
|
||||
assert_eq!(snapshot.active_entries, 2);
|
||||
assert_eq!(snapshot.recent_entries, 2);
|
||||
assert_eq!(snapshot.cleanup_queue_len, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_compact_prunes_stale_recent_entries() {
|
||||
let tracker = UserIpTracker::new();
|
||||
tracker
|
||||
.set_limit_policy(UserMaxUniqueIpsMode::TimeWindow, 1)
|
||||
.await;
|
||||
|
||||
let stale_user = "stale-user".to_string();
|
||||
let stale_ip = test_ipv4(10, 3, 0, 1);
|
||||
{
|
||||
let mut recent_ips = tracker.recent_ips.write().await;
|
||||
recent_ips
|
||||
.entry(stale_user.clone())
|
||||
.or_insert_with(HashMap::new)
|
||||
.insert(stale_ip, Instant::now() - Duration::from_secs(5));
|
||||
}
|
||||
|
||||
tracker.last_compact_epoch_secs.store(0, Ordering::Relaxed);
|
||||
tracker
|
||||
.check_and_add("trigger-user", test_ipv4(10, 3, 0, 2))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let recent_ips = tracker.recent_ips.read().await;
|
||||
let stale_exists = recent_ips
|
||||
.get(&stale_user)
|
||||
.map(|ips| ips.contains_key(&stale_ip))
|
||||
.unwrap_or(false);
|
||||
assert!(!stale_exists);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
102
src/metrics.rs
102
src/metrics.rs
|
|
@ -293,6 +293,27 @@ async fn render_metrics(stats: &Stats, config: &ProxyConfig, ip_tracker: &UserIp
|
|||
}
|
||||
);
|
||||
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_buffer_pool_buffers_total Snapshot of pooled and allocated buffers"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_buffer_pool_buffers_total gauge");
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_buffer_pool_buffers_total{{kind=\"pooled\"}} {}",
|
||||
stats.get_buffer_pool_pooled_gauge()
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_buffer_pool_buffers_total{{kind=\"allocated\"}} {}",
|
||||
stats.get_buffer_pool_allocated_gauge()
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_buffer_pool_buffers_total{{kind=\"in_use\"}} {}",
|
||||
stats.get_buffer_pool_in_use_gauge()
|
||||
);
|
||||
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_connections_total Total accepted connections"
|
||||
|
|
@ -941,6 +962,39 @@ async fn render_metrics(stats: &Stats, config: &ProxyConfig, ip_tracker: &UserIp
|
|||
}
|
||||
);
|
||||
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_me_c2me_enqueue_events_total ME client->ME enqueue outcomes"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_me_c2me_enqueue_events_total counter");
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_c2me_enqueue_events_total{{event=\"full\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_c2me_send_full_total()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_c2me_enqueue_events_total{{event=\"high_water\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_c2me_send_high_water_total()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_me_c2me_enqueue_events_total{{event=\"timeout\"}} {}",
|
||||
if me_allows_normal {
|
||||
stats.get_me_c2me_send_timeout_total()
|
||||
} else {
|
||||
0
|
||||
}
|
||||
);
|
||||
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_me_d2c_batches_total Total DC->Client flush batches"
|
||||
|
|
@ -2490,6 +2544,48 @@ async fn render_metrics(stats: &Stats, config: &ProxyConfig, ip_tracker: &UserIp
|
|||
if user_enabled { 0 } else { 1 }
|
||||
);
|
||||
|
||||
let ip_memory = ip_tracker.memory_stats().await;
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_ip_tracker_users Number of users tracked by IP limiter state"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_ip_tracker_users gauge");
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_ip_tracker_users{{scope=\"active\"}} {}",
|
||||
ip_memory.active_users
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_ip_tracker_users{{scope=\"recent\"}} {}",
|
||||
ip_memory.recent_users
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_ip_tracker_entries Number of IP entries tracked by limiter state"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_ip_tracker_entries gauge");
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_ip_tracker_entries{{scope=\"active\"}} {}",
|
||||
ip_memory.active_entries
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_ip_tracker_entries{{scope=\"recent\"}} {}",
|
||||
ip_memory.recent_entries
|
||||
);
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"# HELP telemt_ip_tracker_cleanup_queue_len Deferred disconnect cleanup queue length"
|
||||
);
|
||||
let _ = writeln!(out, "# TYPE telemt_ip_tracker_cleanup_queue_len gauge");
|
||||
let _ = writeln!(
|
||||
out,
|
||||
"telemt_ip_tracker_cleanup_queue_len {}",
|
||||
ip_memory.cleanup_queue_len
|
||||
);
|
||||
|
||||
if user_enabled {
|
||||
for entry in stats.iter_user_stats() {
|
||||
let user = entry.key();
|
||||
|
|
@ -2728,6 +2824,9 @@ mod tests {
|
|||
assert!(output.contains("telemt_user_unique_ips_recent_window{user=\"alice\"} 1"));
|
||||
assert!(output.contains("telemt_user_unique_ips_limit{user=\"alice\"} 4"));
|
||||
assert!(output.contains("telemt_user_unique_ips_utilization{user=\"alice\"} 0.250000"));
|
||||
assert!(output.contains("telemt_ip_tracker_users{scope=\"active\"} 1"));
|
||||
assert!(output.contains("telemt_ip_tracker_entries{scope=\"active\"} 1"));
|
||||
assert!(output.contains("telemt_ip_tracker_cleanup_queue_len 0"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
@ -2799,6 +2898,9 @@ mod tests {
|
|||
assert!(output.contains("# TYPE telemt_user_unique_ips_recent_window gauge"));
|
||||
assert!(output.contains("# TYPE telemt_user_unique_ips_limit gauge"));
|
||||
assert!(output.contains("# TYPE telemt_user_unique_ips_utilization gauge"));
|
||||
assert!(output.contains("# TYPE telemt_ip_tracker_users gauge"));
|
||||
assert!(output.contains("# TYPE telemt_ip_tracker_entries gauge"));
|
||||
assert!(output.contains("# TYPE telemt_ip_tracker_cleanup_queue_len gauge"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
|
|||
|
|
@ -276,6 +276,7 @@ where
|
|||
stats.increment_user_connects(user);
|
||||
let _direct_connection_lease = stats.acquire_direct_connection_lease();
|
||||
|
||||
let buffer_pool_trim = Arc::clone(&buffer_pool);
|
||||
let relay_result = relay_bidirectional(
|
||||
client_reader,
|
||||
client_writer,
|
||||
|
|
@ -321,6 +322,13 @@ where
|
|||
Err(e) => debug!(user = %user, error = %e, "Direct relay ended with error"),
|
||||
}
|
||||
|
||||
buffer_pool_trim.trim_to(buffer_pool_trim.max_buffers().min(64));
|
||||
let pool_snapshot = buffer_pool_trim.stats();
|
||||
stats.set_buffer_pool_gauges(
|
||||
pool_snapshot.pooled,
|
||||
pool_snapshot.allocated,
|
||||
pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled),
|
||||
);
|
||||
relay_result
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -645,11 +645,14 @@ async fn enqueue_c2me_command(
|
|||
tx: &mpsc::Sender<C2MeCommand>,
|
||||
cmd: C2MeCommand,
|
||||
send_timeout: Option<Duration>,
|
||||
stats: &Stats,
|
||||
) -> std::result::Result<(), mpsc::error::SendError<C2MeCommand>> {
|
||||
match tx.try_send(cmd) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(mpsc::error::TrySendError::Closed(cmd)) => Err(mpsc::error::SendError(cmd)),
|
||||
Err(mpsc::error::TrySendError::Full(cmd)) => {
|
||||
stats.increment_me_c2me_send_full_total();
|
||||
stats.increment_me_c2me_send_high_water_total();
|
||||
note_relay_pressure_event();
|
||||
// Cooperative yield reduces burst catch-up when the per-conn queue is near saturation.
|
||||
if tx.capacity() <= C2ME_SOFT_PRESSURE_MIN_FREE_SLOTS {
|
||||
|
|
@ -658,7 +661,10 @@ async fn enqueue_c2me_command(
|
|||
let reserve_result = match send_timeout {
|
||||
Some(send_timeout) => match timeout(send_timeout, tx.reserve()).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => return Err(mpsc::error::SendError(cmd)),
|
||||
Err(_) => {
|
||||
stats.increment_me_c2me_send_timeout_total();
|
||||
return Err(mpsc::error::SendError(cmd));
|
||||
}
|
||||
},
|
||||
None => tx.reserve().await,
|
||||
};
|
||||
|
|
@ -667,7 +673,10 @@ async fn enqueue_c2me_command(
|
|||
permit.send(cmd);
|
||||
Ok(())
|
||||
}
|
||||
Err(_) => Err(mpsc::error::SendError(cmd)),
|
||||
Err(_) => {
|
||||
stats.increment_me_c2me_send_timeout_total();
|
||||
Err(mpsc::error::SendError(cmd))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -841,11 +850,23 @@ where
|
|||
let me_writer = tokio::spawn(async move {
|
||||
let mut writer = crypto_writer;
|
||||
let mut frame_buf = Vec::with_capacity(16 * 1024);
|
||||
let shrink_threshold = d2c_flush_policy.frame_buf_shrink_threshold_bytes;
|
||||
|
||||
fn shrink_session_vec(buf: &mut Vec<u8>, threshold: usize) {
|
||||
if buf.capacity() > threshold {
|
||||
buf.clear();
|
||||
buf.shrink_to(threshold);
|
||||
} else {
|
||||
buf.clear();
|
||||
}
|
||||
}
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
msg = me_rx_task.recv() => {
|
||||
let Some(first) = msg else {
|
||||
debug!(conn_id, "ME channel closed");
|
||||
shrink_session_vec(&mut frame_buf, shrink_threshold);
|
||||
return Err(ProxyError::Proxy("ME connection lost".into()));
|
||||
};
|
||||
|
||||
|
|
@ -901,6 +922,7 @@ where
|
|||
batch_bytes,
|
||||
flush_duration_us,
|
||||
);
|
||||
shrink_session_vec(&mut frame_buf, shrink_threshold);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
|
@ -962,6 +984,7 @@ where
|
|||
batch_bytes,
|
||||
flush_duration_us,
|
||||
);
|
||||
shrink_session_vec(&mut frame_buf, shrink_threshold);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
|
@ -1027,6 +1050,7 @@ where
|
|||
batch_bytes,
|
||||
flush_duration_us,
|
||||
);
|
||||
shrink_session_vec(&mut frame_buf, shrink_threshold);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
|
@ -1091,6 +1115,7 @@ where
|
|||
batch_bytes,
|
||||
flush_duration_us,
|
||||
);
|
||||
shrink_session_vec(&mut frame_buf, shrink_threshold);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
|
@ -1098,6 +1123,7 @@ where
|
|||
}
|
||||
Ok(None) => {
|
||||
debug!(conn_id, "ME channel closed");
|
||||
shrink_session_vec(&mut frame_buf, shrink_threshold);
|
||||
return Err(ProxyError::Proxy("ME connection lost".into()));
|
||||
}
|
||||
Err(_) => {
|
||||
|
|
@ -1147,6 +1173,7 @@ where
|
|||
}
|
||||
_ = &mut stop_rx => {
|
||||
debug!(conn_id, "ME writer stop signal");
|
||||
shrink_session_vec(&mut frame_buf, shrink_threshold);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
|
@ -1172,7 +1199,13 @@ where
|
|||
user = %user,
|
||||
"Middle-relay pressure eviction for idle-candidate session"
|
||||
);
|
||||
let _ = enqueue_c2me_command(&c2me_tx, C2MeCommand::Close, c2me_send_timeout).await;
|
||||
let _ = enqueue_c2me_command(
|
||||
&c2me_tx,
|
||||
C2MeCommand::Close,
|
||||
c2me_send_timeout,
|
||||
stats.as_ref(),
|
||||
)
|
||||
.await;
|
||||
main_result = Err(ProxyError::Proxy(
|
||||
"middle-relay session evicted under pressure (idle-candidate)".to_string(),
|
||||
));
|
||||
|
|
@ -1191,7 +1224,13 @@ where
|
|||
"Cutover affected middle session, closing client connection"
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
let _ = enqueue_c2me_command(&c2me_tx, C2MeCommand::Close, c2me_send_timeout).await;
|
||||
let _ = enqueue_c2me_command(
|
||||
&c2me_tx,
|
||||
C2MeCommand::Close,
|
||||
c2me_send_timeout,
|
||||
stats.as_ref(),
|
||||
)
|
||||
.await;
|
||||
main_result = Err(ProxyError::Proxy(ROUTE_SWITCH_ERROR_MSG.to_string()));
|
||||
break;
|
||||
}
|
||||
|
|
@ -1253,6 +1292,7 @@ where
|
|||
&c2me_tx,
|
||||
C2MeCommand::Data { payload, flags },
|
||||
c2me_send_timeout,
|
||||
stats.as_ref(),
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
|
|
@ -1264,8 +1304,12 @@ where
|
|||
Ok(None) => {
|
||||
debug!(conn_id, "Client EOF");
|
||||
client_closed = true;
|
||||
let _ =
|
||||
enqueue_c2me_command(&c2me_tx, C2MeCommand::Close, c2me_send_timeout)
|
||||
let _ = enqueue_c2me_command(
|
||||
&c2me_tx,
|
||||
C2MeCommand::Close,
|
||||
c2me_send_timeout,
|
||||
stats.as_ref(),
|
||||
)
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
|
|
@ -1317,6 +1361,13 @@ where
|
|||
);
|
||||
clear_relay_idle_candidate(conn_id);
|
||||
me_pool.registry().unregister(conn_id).await;
|
||||
buffer_pool.trim_to(buffer_pool.max_buffers().min(64));
|
||||
let pool_snapshot = buffer_pool.stats();
|
||||
stats.set_buffer_pool_gauges(
|
||||
pool_snapshot.pooled,
|
||||
pool_snapshot.allocated,
|
||||
pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled),
|
||||
);
|
||||
result
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use super::*;
|
||||
use crate::stats::Stats;
|
||||
use crate::stream::BufferPool;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
|
@ -119,6 +120,7 @@ async fn c2me_channel_full_path_yields_then_sends() {
|
|||
.expect("priming queue with one frame must succeed");
|
||||
|
||||
let tx2 = tx.clone();
|
||||
let stats = Stats::default();
|
||||
let producer = tokio::spawn(async move {
|
||||
enqueue_c2me_command(
|
||||
&tx2,
|
||||
|
|
@ -127,6 +129,7 @@ async fn c2me_channel_full_path_yields_then_sends() {
|
|||
flags: 2,
|
||||
},
|
||||
None,
|
||||
&stats,
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
|
|
|||
|
|
@ -200,6 +200,14 @@ pub struct Stats {
|
|||
me_d2c_flush_duration_us_bucket_1001_5000: AtomicU64,
|
||||
me_d2c_flush_duration_us_bucket_5001_20000: AtomicU64,
|
||||
me_d2c_flush_duration_us_bucket_gt_20000: AtomicU64,
|
||||
// Buffer pool gauges
|
||||
buffer_pool_pooled_gauge: AtomicU64,
|
||||
buffer_pool_allocated_gauge: AtomicU64,
|
||||
buffer_pool_in_use_gauge: AtomicU64,
|
||||
// C2ME enqueue observability
|
||||
me_c2me_send_full_total: AtomicU64,
|
||||
me_c2me_send_high_water_total: AtomicU64,
|
||||
me_c2me_send_timeout_total: AtomicU64,
|
||||
me_d2c_batch_timeout_armed_total: AtomicU64,
|
||||
me_d2c_batch_timeout_fired_total: AtomicU64,
|
||||
me_writer_pick_sorted_rr_success_try_total: AtomicU64,
|
||||
|
|
@ -1414,6 +1422,37 @@ impl Stats {
|
|||
.store(value, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_buffer_pool_gauges(&self, pooled: usize, allocated: usize, in_use: usize) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.buffer_pool_pooled_gauge
|
||||
.store(pooled as u64, Ordering::Relaxed);
|
||||
self.buffer_pool_allocated_gauge
|
||||
.store(allocated as u64, Ordering::Relaxed);
|
||||
self.buffer_pool_in_use_gauge
|
||||
.store(in_use as u64, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn increment_me_c2me_send_full_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_c2me_send_full_total.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn increment_me_c2me_send_high_water_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_c2me_send_high_water_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn increment_me_c2me_send_timeout_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_c2me_send_timeout_total
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
pub fn increment_me_floor_cap_block_total(&self) {
|
||||
if self.telemetry_me_allows_normal() {
|
||||
self.me_floor_cap_block_total
|
||||
|
|
@ -1780,6 +1819,30 @@ impl Stats {
|
|||
self.me_d2c_flush_duration_us_bucket_gt_20000
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn get_buffer_pool_pooled_gauge(&self) -> u64 {
|
||||
self.buffer_pool_pooled_gauge.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn get_buffer_pool_allocated_gauge(&self) -> u64 {
|
||||
self.buffer_pool_allocated_gauge.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn get_buffer_pool_in_use_gauge(&self) -> u64 {
|
||||
self.buffer_pool_in_use_gauge.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn get_me_c2me_send_full_total(&self) -> u64 {
|
||||
self.me_c2me_send_full_total.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn get_me_c2me_send_high_water_total(&self) -> u64 {
|
||||
self.me_c2me_send_high_water_total.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn get_me_c2me_send_timeout_total(&self) -> u64 {
|
||||
self.me_c2me_send_timeout_total.load(Ordering::Relaxed)
|
||||
}
|
||||
pub fn get_me_d2c_batch_timeout_armed_total(&self) -> u64 {
|
||||
self.me_d2c_batch_timeout_armed_total
|
||||
.load(Ordering::Relaxed)
|
||||
|
|
@ -2171,6 +2234,8 @@ impl ReplayShard {
|
|||
|
||||
fn cleanup(&mut self, now: Instant, window: Duration) {
|
||||
if window.is_zero() {
|
||||
self.cache.clear();
|
||||
self.queue.clear();
|
||||
return;
|
||||
}
|
||||
let cutoff = now.checked_sub(window).unwrap_or(now);
|
||||
|
|
@ -2192,13 +2257,22 @@ impl ReplayShard {
|
|||
}
|
||||
|
||||
fn check(&mut self, key: &[u8], now: Instant, window: Duration) -> bool {
|
||||
if window.is_zero() {
|
||||
return false;
|
||||
}
|
||||
self.cleanup(now, window);
|
||||
// key is &[u8], resolves Q=[u8] via Box<[u8]>: Borrow<[u8]>
|
||||
self.cache.get(key).is_some()
|
||||
}
|
||||
|
||||
fn add(&mut self, key: &[u8], now: Instant, window: Duration) {
|
||||
if window.is_zero() {
|
||||
return;
|
||||
}
|
||||
self.cleanup(now, window);
|
||||
if self.cache.peek(key).is_some() {
|
||||
return;
|
||||
}
|
||||
|
||||
let seq = self.next_seq();
|
||||
let boxed_key: Box<[u8]> = key.into();
|
||||
|
|
@ -2341,7 +2415,7 @@ impl ReplayChecker {
|
|||
let interval = if self.window.as_secs() > 60 {
|
||||
Duration::from_secs(30)
|
||||
} else {
|
||||
Duration::from_secs(self.window.as_secs().max(1) / 2)
|
||||
Duration::from_secs((self.window.as_secs().max(1) / 2).max(1))
|
||||
};
|
||||
|
||||
loop {
|
||||
|
|
@ -2553,6 +2627,20 @@ mod tests {
|
|||
assert!(!checker.check_handshake(b"expire"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_replay_checker_zero_window_does_not_retain_entries() {
|
||||
let checker = ReplayChecker::new(100, Duration::ZERO);
|
||||
|
||||
for _ in 0..1_000 {
|
||||
assert!(!checker.check_handshake(b"no-retain"));
|
||||
checker.add_handshake(b"no-retain");
|
||||
}
|
||||
|
||||
let stats = checker.stats();
|
||||
assert_eq!(stats.total_entries, 0);
|
||||
assert_eq!(stats.total_queue_len, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_replay_checker_stats() {
|
||||
let checker = ReplayChecker::new(100, Duration::from_secs(60));
|
||||
|
|
|
|||
|
|
@ -35,6 +35,10 @@ pub struct BufferPool {
|
|||
misses: AtomicUsize,
|
||||
/// Number of successful reuses
|
||||
hits: AtomicUsize,
|
||||
/// Number of non-standard buffers replaced with a fresh default-sized buffer
|
||||
replaced_nonstandard: AtomicUsize,
|
||||
/// Number of buffers dropped because the pool queue was full
|
||||
dropped_pool_full: AtomicUsize,
|
||||
}
|
||||
|
||||
impl BufferPool {
|
||||
|
|
@ -52,6 +56,8 @@ impl BufferPool {
|
|||
allocated: AtomicUsize::new(0),
|
||||
misses: AtomicUsize::new(0),
|
||||
hits: AtomicUsize::new(0),
|
||||
replaced_nonstandard: AtomicUsize::new(0),
|
||||
dropped_pool_full: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -91,17 +97,36 @@ impl BufferPool {
|
|||
|
||||
/// Return a buffer to the pool
|
||||
fn return_buffer(&self, mut buffer: BytesMut) {
|
||||
// Clear the buffer but keep capacity
|
||||
buffer.clear();
|
||||
const MAX_RETAINED_BUFFER_FACTOR: usize = 2;
|
||||
|
||||
// Only return if we haven't exceeded max and buffer is right size
|
||||
if buffer.capacity() >= self.buffer_size {
|
||||
// Try to push to pool, if full just drop
|
||||
let _ = self.buffers.push(buffer);
|
||||
// Clear the buffer but keep capacity.
|
||||
buffer.clear();
|
||||
let max_retained_capacity = self
|
||||
.buffer_size
|
||||
.saturating_mul(MAX_RETAINED_BUFFER_FACTOR)
|
||||
.max(self.buffer_size);
|
||||
|
||||
// Keep only near-default capacities in the pool. Oversized buffers keep
|
||||
// RSS elevated for hours under churn; replace them with default-sized
|
||||
// buffers before re-pooling.
|
||||
if buffer.capacity() < self.buffer_size || buffer.capacity() > max_retained_capacity {
|
||||
self.replaced_nonstandard.fetch_add(1, Ordering::Relaxed);
|
||||
buffer = BytesMut::with_capacity(self.buffer_size);
|
||||
}
|
||||
// If buffer was dropped (pool full), decrement allocated
|
||||
// Actually we don't decrement here because the buffer might have been
|
||||
// grown beyond our size - we just let it go
|
||||
|
||||
// Try to return into the queue; if full, drop and update accounting.
|
||||
if self.buffers.push(buffer).is_err() {
|
||||
self.dropped_pool_full.fetch_add(1, Ordering::Relaxed);
|
||||
self.decrement_allocated();
|
||||
}
|
||||
}
|
||||
|
||||
fn decrement_allocated(&self) {
|
||||
let _ = self
|
||||
.allocated
|
||||
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
|
||||
Some(current.saturating_sub(1))
|
||||
});
|
||||
}
|
||||
|
||||
/// Get pool statistics
|
||||
|
|
@ -113,6 +138,8 @@ impl BufferPool {
|
|||
buffer_size: self.buffer_size,
|
||||
hits: self.hits.load(Ordering::Relaxed),
|
||||
misses: self.misses.load(Ordering::Relaxed),
|
||||
replaced_nonstandard: self.replaced_nonstandard.load(Ordering::Relaxed),
|
||||
dropped_pool_full: self.dropped_pool_full.load(Ordering::Relaxed),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -121,6 +148,41 @@ impl BufferPool {
|
|||
self.buffer_size
|
||||
}
|
||||
|
||||
/// Maximum number of buffers the pool will retain.
|
||||
pub fn max_buffers(&self) -> usize {
|
||||
self.max_buffers
|
||||
}
|
||||
|
||||
/// Current number of pooled buffers.
|
||||
pub fn pooled(&self) -> usize {
|
||||
self.buffers.len()
|
||||
}
|
||||
|
||||
/// Total buffers allocated (pooled + checked out).
|
||||
pub fn allocated(&self) -> usize {
|
||||
self.allocated.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Best-effort number of buffers currently checked out.
|
||||
pub fn in_use(&self) -> usize {
|
||||
self.allocated().saturating_sub(self.pooled())
|
||||
}
|
||||
|
||||
/// Trim pooled buffers down to a target count.
|
||||
pub fn trim_to(&self, target_pooled: usize) {
|
||||
let target = target_pooled.min(self.max_buffers);
|
||||
loop {
|
||||
if self.buffers.len() <= target {
|
||||
break;
|
||||
}
|
||||
if self.buffers.pop().is_some() {
|
||||
self.decrement_allocated();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Preallocate buffers to fill the pool
|
||||
pub fn preallocate(&self, count: usize) {
|
||||
let to_alloc = count.min(self.max_buffers);
|
||||
|
|
@ -160,6 +222,10 @@ pub struct PoolStats {
|
|||
pub hits: usize,
|
||||
/// Number of cache misses (new allocation)
|
||||
pub misses: usize,
|
||||
/// Number of non-standard buffers replaced during return
|
||||
pub replaced_nonstandard: usize,
|
||||
/// Number of buffers dropped because the pool queue was full
|
||||
pub dropped_pool_full: usize,
|
||||
}
|
||||
|
||||
impl PoolStats {
|
||||
|
|
@ -185,6 +251,7 @@ pub struct PooledBuffer {
|
|||
impl PooledBuffer {
|
||||
/// Take the inner buffer, preventing return to pool
|
||||
pub fn take(mut self) -> BytesMut {
|
||||
self.pool.decrement_allocated();
|
||||
self.buffer.take().unwrap()
|
||||
}
|
||||
|
||||
|
|
@ -364,6 +431,25 @@ mod tests {
|
|||
|
||||
let stats = pool.stats();
|
||||
assert_eq!(stats.pooled, 0);
|
||||
assert_eq!(stats.allocated, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pool_replaces_oversized_buffers() {
|
||||
let pool = Arc::new(BufferPool::with_config(1024, 10));
|
||||
|
||||
{
|
||||
let mut buf = pool.get();
|
||||
buf.reserve(8192);
|
||||
assert!(buf.capacity() > 2048);
|
||||
}
|
||||
|
||||
let stats = pool.stats();
|
||||
assert_eq!(stats.replaced_nonstandard, 1);
|
||||
assert_eq!(stats.pooled, 1);
|
||||
|
||||
let buf = pool.get();
|
||||
assert!(buf.capacity() <= 2048);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
Loading…
Reference in New Issue