Bound ME writer queues by resident payload bytes

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-07-11 18:43:42 +03:00
parent 893ce0cf36
commit d4c4980e5a
25 changed files with 798 additions and 82 deletions
+16
View File
@@ -24,6 +24,10 @@ const DEFAULT_ME_ADAPTIVE_FLOOR_MAX_WARM_WRITERS_GLOBAL: u32 = 256;
const DEFAULT_ME_ROUTE_BACKPRESSURE_ENABLED: bool = false; const DEFAULT_ME_ROUTE_BACKPRESSURE_ENABLED: bool = false;
const DEFAULT_ME_ROUTE_FAIRSHARE_ENABLED: bool = false; const DEFAULT_ME_ROUTE_FAIRSHARE_ENABLED: bool = false;
const DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 4096; const DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 4096;
pub(crate) const ME_WRITER_BYTE_PERMIT_UNIT_BYTES: usize = 16 * 1024;
pub(crate) const ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES: usize = 256;
const DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES: usize =
32 * 1024 * 1024 + ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
const DEFAULT_ME_ROUTE_CHANNEL_CAPACITY: usize = 768; const DEFAULT_ME_ROUTE_CHANNEL_CAPACITY: usize = 768;
const DEFAULT_ME_C2ME_CHANNEL_CAPACITY: usize = 1024; const DEFAULT_ME_C2ME_CHANNEL_CAPACITY: usize = 1024;
const DEFAULT_ME_READER_ROUTE_DATA_WAIT_MS: u64 = 2; const DEFAULT_ME_READER_ROUTE_DATA_WAIT_MS: u64 = 2;
@@ -455,6 +459,18 @@ pub(crate) fn default_me_writer_cmd_channel_capacity() -> usize {
DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY
} }
pub(crate) fn default_me_writer_byte_budget_bytes() -> usize {
DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES
}
pub(crate) fn minimum_me_writer_byte_budget_bytes(max_client_frame: usize) -> usize {
max_client_frame
.saturating_mul(2)
.saturating_add(ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES)
.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
.saturating_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
}
pub(crate) fn default_me_route_channel_capacity() -> usize { pub(crate) fn default_me_route_channel_capacity() -> usize {
DEFAULT_ME_ROUTE_CHANNEL_CAPACITY DEFAULT_ME_ROUTE_CHANNEL_CAPACITY
} }
+17
View File
@@ -39,6 +39,7 @@ use self::validation::{
}; };
const MAX_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 16_384; const MAX_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 16_384;
const MAX_ME_WRITER_BYTE_BUDGET_BYTES: usize = 256 * 1024 * 1024;
const MAX_ME_ROUTE_CHANNEL_CAPACITY: usize = 8_192; const MAX_ME_ROUTE_CHANNEL_CAPACITY: usize = 8_192;
const MAX_ME_C2ME_CHANNEL_CAPACITY: usize = 8_192; const MAX_ME_C2ME_CHANNEL_CAPACITY: usize = 8_192;
const MIN_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024; const MIN_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024;
@@ -540,6 +541,22 @@ impl ProxyConfig {
))); )));
} }
let min_writer_byte_budget =
minimum_me_writer_byte_budget_bytes(config.general.max_client_frame);
if config.general.me_writer_byte_budget_bytes % ME_WRITER_BYTE_PERMIT_UNIT_BYTES != 0 {
return Err(ProxyError::Config(format!(
"general.me_writer_byte_budget_bytes must be a multiple of {ME_WRITER_BYTE_PERMIT_UNIT_BYTES}"
)));
}
if !(min_writer_byte_budget..=MAX_ME_WRITER_BYTE_BUDGET_BYTES)
.contains(&config.general.me_writer_byte_budget_bytes)
{
return Err(ProxyError::Config(format!(
"general.me_writer_byte_budget_bytes must be within [{min_writer_byte_budget}, {MAX_ME_WRITER_BYTE_BUDGET_BYTES}] for general.max_client_frame={}",
config.general.max_client_frame
)));
}
if config.general.me_c2me_send_timeout_ms > 60_000 { if config.general.me_c2me_send_timeout_ms > 60_000 {
return Err(ProxyError::Config( return Err(ProxyError::Config(
"general.me_c2me_send_timeout_ms must be within [0, 60000]".to_string(), "general.me_c2me_send_timeout_ms must be within [0, 60000]".to_string(),
+1
View File
@@ -52,6 +52,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
"me_keepalive_payload_random", "me_keepalive_payload_random",
"rpc_proxy_req_every", "rpc_proxy_req_every",
"me_writer_cmd_channel_capacity", "me_writer_cmd_channel_capacity",
"me_writer_byte_budget_bytes",
"me_route_channel_capacity", "me_route_channel_capacity",
"me_c2me_channel_capacity", "me_c2me_channel_capacity",
"me_c2me_send_timeout_ms", "me_c2me_send_timeout_ms",
@@ -95,6 +95,67 @@ max_client_frame = 16777217
remove_temp_config(&path); remove_temp_config(&path);
} }
#[test]
fn load_rejects_writer_byte_budget_below_frame_residency_minimum() {
let path = write_temp_config(
r#"
[general]
max_client_frame = 16777216
me_writer_byte_budget_bytes = 33554432
"#,
);
let err = ProxyConfig::load(&path)
.expect_err("writer byte budget below frame residency minimum must fail");
let msg = err.to_string();
assert!(
msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"),
"error must explain writer byte budget minimum, got: {msg}"
);
remove_temp_config(&path);
}
#[test]
fn load_rejects_unaligned_writer_byte_budget() {
let path = write_temp_config(
r#"
[general]
me_writer_byte_budget_bytes = 33570817
"#,
);
let err = ProxyConfig::load(&path)
.expect_err("writer byte budget outside permit granularity must fail");
let msg = err.to_string();
assert!(
msg.contains("general.me_writer_byte_budget_bytes must be a multiple of 16384"),
"error must explain writer byte budget alignment, got: {msg}"
);
remove_temp_config(&path);
}
#[test]
fn load_rejects_writer_byte_budget_above_hard_cap() {
let path = write_temp_config(
r#"
[general]
me_writer_byte_budget_bytes = 268451840
"#,
);
let err =
ProxyConfig::load(&path).expect_err("writer byte budget above hard cap must fail");
let msg = err.to_string();
assert!(
msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"),
"error must explain writer byte budget hard cap, got: {msg}"
);
remove_temp_config(&path);
}
#[test] #[test]
fn load_rejects_listen_backlog_above_i32_upper_bound() { fn load_rejects_listen_backlog_above_i32_upper_bound() {
let path = write_temp_config( let path = write_temp_config(
@@ -139,6 +200,7 @@ fn load_accepts_memory_limits_at_hard_upper_bounds() {
r#" r#"
[general] [general]
me_writer_cmd_channel_capacity = 16384 me_writer_cmd_channel_capacity = 16384
me_writer_byte_budget_bytes = 268435456
me_route_channel_capacity = 8192 me_route_channel_capacity = 8192
me_c2me_channel_capacity = 8192 me_c2me_channel_capacity = 8192
max_client_frame = 16777216 max_client_frame = 16777216
@@ -147,6 +209,7 @@ max_client_frame = 16777216
let cfg = ProxyConfig::load(&path).expect("hard upper bound values must be accepted"); let cfg = ProxyConfig::load(&path).expect("hard upper bound values must be accepted");
assert_eq!(cfg.general.me_writer_cmd_channel_capacity, 16384); assert_eq!(cfg.general.me_writer_cmd_channel_capacity, 16384);
assert_eq!(cfg.general.me_writer_byte_budget_bytes, 256 * 1024 * 1024);
assert_eq!(cfg.general.me_route_channel_capacity, 8192); assert_eq!(cfg.general.me_route_channel_capacity, 8192);
assert_eq!(cfg.general.me_c2me_channel_capacity, 8192); assert_eq!(cfg.general.me_c2me_channel_capacity, 8192);
assert_eq!(cfg.general.max_client_frame, 16 * 1024 * 1024); assert_eq!(cfg.general.max_client_frame, 16 * 1024 * 1024);
+5
View File
@@ -579,6 +579,10 @@ pub struct GeneralConfig {
#[serde(default = "default_me_writer_cmd_channel_capacity")] #[serde(default = "default_me_writer_cmd_channel_capacity")]
pub me_writer_cmd_channel_capacity: usize, pub me_writer_cmd_channel_capacity: usize,
/// Resident-memory budget in bytes for each ME writer data queue.
#[serde(default = "default_me_writer_byte_budget_bytes")]
pub me_writer_byte_budget_bytes: usize,
/// Capacity of per-connection ME response route channel. /// Capacity of per-connection ME response route channel.
#[serde(default = "default_me_route_channel_capacity")] #[serde(default = "default_me_route_channel_capacity")]
pub me_route_channel_capacity: usize, pub me_route_channel_capacity: usize,
@@ -1103,6 +1107,7 @@ impl Default for GeneralConfig {
me_keepalive_payload_random: default_true(), me_keepalive_payload_random: default_true(),
rpc_proxy_req_every: default_rpc_proxy_req_every(), rpc_proxy_req_every: default_rpc_proxy_req_every(),
me_writer_cmd_channel_capacity: default_me_writer_cmd_channel_capacity(), me_writer_cmd_channel_capacity: default_me_writer_cmd_channel_capacity(),
me_writer_byte_budget_bytes: default_me_writer_byte_budget_bytes(),
me_route_channel_capacity: default_me_route_channel_capacity(), me_route_channel_capacity: default_me_route_channel_capacity(),
me_c2me_channel_capacity: default_me_c2me_channel_capacity(), me_c2me_channel_capacity: default_me_c2me_channel_capacity(),
me_c2me_send_timeout_ms: default_me_c2me_send_timeout_ms(), me_c2me_send_timeout_ms: default_me_c2me_send_timeout_ms(),
+1
View File
@@ -279,6 +279,7 @@ pub(crate) async fn initialize_me_pool(
config.general.me_writer_pick_sample_size, config.general.me_writer_pick_sample_size,
config.general.me_socks_kdf_policy, config.general.me_socks_kdf_policy,
config.general.me_writer_cmd_channel_capacity, config.general.me_writer_cmd_channel_capacity,
config.general.me_writer_byte_budget_bytes,
config.general.me_route_channel_capacity, config.general.me_route_channel_capacity,
config.general.me_route_backpressure_enabled, config.general.me_route_backpressure_enabled,
config.general.me_route_fairshare_enabled, config.general.me_route_fairshare_enabled,
+79
View File
@@ -2435,6 +2435,85 @@ async fn render_metrics(
} }
); );
let _ = writeln!(
out,
"# HELP telemt_me_writer_byte_budget_limit_bytes Configured resident-memory budget per ME writer"
);
let _ = writeln!(
out,
"# TYPE telemt_me_writer_byte_budget_limit_bytes gauge"
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_limit_bytes {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_limit_bytes_gauge()
} else {
0
}
);
let _ = writeln!(
out,
"# HELP telemt_me_writer_byte_budget_reserved_bytes Aggregate ME writer memory reservations by lifecycle state"
);
let _ = writeln!(
out,
"# TYPE telemt_me_writer_byte_budget_reserved_bytes gauge"
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_reserved_bytes{{state=\"queued\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_queued_bytes_gauge()
} else {
0
}
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_reserved_bytes{{state=\"inflight\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_inflight_bytes_gauge()
} else {
0
}
);
let _ = writeln!(
out,
"# HELP telemt_me_writer_byte_budget_events_total ME writer byte-budget outcomes"
);
let _ = writeln!(
out,
"# TYPE telemt_me_writer_byte_budget_events_total counter"
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_events_total{{result=\"wait\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_wait_total()
} else {
0
}
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_events_total{{result=\"timeout\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_timeout_total()
} else {
0
}
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_events_total{{result=\"oversize\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_oversize_total()
} else {
0
}
);
let _ = writeln!( let _ = writeln!(
out, out,
"# HELP telemt_me_writer_pick_total ME writer-pick outcomes by mode and result" "# HELP telemt_me_writer_pick_total ME writer-pick outcomes by mode and result"
+29
View File
@@ -269,6 +269,35 @@ impl Stats {
self.me_writer_pick_mode_switch_total self.me_writer_pick_mode_switch_total
.load(Ordering::Relaxed) .load(Ordering::Relaxed)
} }
/// Returns the configured resident-memory limit per ME writer.
pub fn get_me_writer_byte_budget_limit_bytes_gauge(&self) -> u64 {
self.me_writer_byte_budget_limit_bytes_gauge
.load(Ordering::Relaxed)
}
/// Returns aggregate queued or enqueueing writer memory reservations.
pub fn get_me_writer_byte_budget_queued_bytes_gauge(&self) -> u64 {
self.me_writer_byte_budget_queued_bytes_gauge
.load(Ordering::Relaxed)
}
/// Returns aggregate writer reservations currently owned by socket writes.
pub fn get_me_writer_byte_budget_inflight_bytes_gauge(&self) -> u64 {
self.me_writer_byte_budget_inflight_bytes_gauge
.load(Ordering::Relaxed)
}
/// Returns the count of blocking writer byte-budget waits.
pub fn get_me_writer_byte_budget_wait_total(&self) -> u64 {
self.me_writer_byte_budget_wait_total.load(Ordering::Relaxed)
}
/// Returns the count of writer byte-budget wait timeouts.
pub fn get_me_writer_byte_budget_timeout_total(&self) -> u64 {
self.me_writer_byte_budget_timeout_total
.load(Ordering::Relaxed)
}
/// Returns the count of payloads that cannot fit the configured writer budget.
pub fn get_me_writer_byte_budget_oversize_total(&self) -> u64 {
self.me_writer_byte_budget_oversize_total
.load(Ordering::Relaxed)
}
pub fn get_me_socks_kdf_strict_reject(&self) -> u64 { pub fn get_me_socks_kdf_strict_reject(&self) -> u64 {
self.me_socks_kdf_strict_reject.load(Ordering::Relaxed) self.me_socks_kdf_strict_reject.load(Ordering::Relaxed)
} }
+6
View File
@@ -292,6 +292,12 @@ pub struct Stats {
me_writer_pick_p2c_no_candidate_total: AtomicU64, me_writer_pick_p2c_no_candidate_total: AtomicU64,
me_writer_pick_blocking_fallback_total: AtomicU64, me_writer_pick_blocking_fallback_total: AtomicU64,
me_writer_pick_mode_switch_total: AtomicU64, me_writer_pick_mode_switch_total: AtomicU64,
me_writer_byte_budget_limit_bytes_gauge: AtomicU64,
me_writer_byte_budget_queued_bytes_gauge: AtomicU64,
me_writer_byte_budget_inflight_bytes_gauge: AtomicU64,
me_writer_byte_budget_wait_total: AtomicU64,
me_writer_byte_budget_timeout_total: AtomicU64,
me_writer_byte_budget_oversize_total: AtomicU64,
me_socks_kdf_strict_reject: AtomicU64, me_socks_kdf_strict_reject: AtomicU64,
me_socks_kdf_compat_fallback: AtomicU64, me_socks_kdf_compat_fallback: AtomicU64,
secure_padding_invalid: AtomicU64, secure_padding_invalid: AtomicU64,
+50
View File
@@ -88,6 +88,56 @@ impl Stats {
.fetch_add(1, Ordering::Relaxed); .fetch_add(1, Ordering::Relaxed);
} }
} }
/// Publishes the configured resident-memory limit applied to each ME writer.
pub fn set_me_writer_byte_budget_limit_bytes(&self, bytes: usize) {
self.me_writer_byte_budget_limit_bytes_gauge
.store(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn add_me_writer_byte_budget_queued_bytes(&self, bytes: u64) {
self.me_writer_byte_budget_queued_bytes_gauge
.fetch_add(bytes, Ordering::Relaxed);
}
pub(crate) fn move_me_writer_byte_budget_to_inflight(&self, bytes: u64) {
let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|current| Some(current.saturating_sub(bytes)),
);
self.me_writer_byte_budget_inflight_bytes_gauge
.fetch_add(bytes, Ordering::Relaxed);
}
pub(crate) fn release_me_writer_byte_budget_queued_bytes(&self, bytes: u64) {
let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|current| Some(current.saturating_sub(bytes)),
);
}
pub(crate) fn release_me_writer_byte_budget_inflight_bytes(&self, bytes: u64) {
let _ = self.me_writer_byte_budget_inflight_bytes_gauge.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|current| Some(current.saturating_sub(bytes)),
);
}
pub(crate) fn increment_me_writer_byte_budget_wait_total(&self) {
if self.telemetry_me_allows_normal() {
self.me_writer_byte_budget_wait_total
.fetch_add(1, Ordering::Relaxed);
}
}
pub(crate) fn increment_me_writer_byte_budget_timeout_total(&self) {
if self.telemetry_me_allows_normal() {
self.me_writer_byte_budget_timeout_total
.fetch_add(1, Ordering::Relaxed);
}
}
pub(crate) fn increment_me_writer_byte_budget_oversize_total(&self) {
if self.telemetry_me_allows_normal() {
self.me_writer_byte_budget_oversize_total
.fetch_add(1, Ordering::Relaxed);
}
}
pub fn increment_me_socks_kdf_strict_reject(&self) { pub fn increment_me_socks_kdf_strict_reject(&self) {
if self.telemetry_me_allows_normal() { if self.telemetry_me_allows_normal() {
self.me_socks_kdf_strict_reject self.me_socks_kdf_strict_reject
+91 -2
View File
@@ -1,10 +1,12 @@
use bytes::Bytes; use bytes::Bytes;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::OwnedSemaphorePermit; use tokio::sync::OwnedSemaphorePermit;
use crate::crypto::{AesCbc, crc32, crc32c}; use crate::crypto::{AesCbc, crc32, crc32c};
use crate::error::{ProxyError, Result}; use crate::error::{ProxyError, Result};
use crate::protocol::constants::*; use crate::protocol::constants::*;
use crate::stats::Stats;
use crate::stream::PooledBuffer; use crate::stream::PooledBuffer;
use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len}; use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len};
@@ -12,11 +14,65 @@ use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len};
const RPC_WRITER_FRAME_BUF_SHRINK_THRESHOLD: usize = 256 * 1024; const RPC_WRITER_FRAME_BUF_SHRINK_THRESHOLD: usize = 256 * 1024;
const RPC_WRITER_FRAME_BUF_RETAIN: usize = 64 * 1024; const RPC_WRITER_FRAME_BUF_RETAIN: usize = 64 * 1024;
enum WriterBytePermitState {
Queued,
Inflight,
}
/// Holds one writer memory reservation through queueing and socket write completion.
pub(crate) struct WriterBytePermit {
_permit: OwnedSemaphorePermit,
reserved_bytes: u64,
stats: Arc<Stats>,
state: WriterBytePermitState,
}
impl WriterBytePermit {
/// Creates a queued reservation and publishes its rounded resident-byte cost.
pub(crate) fn new(
permit: OwnedSemaphorePermit,
reserved_bytes: usize,
stats: Arc<Stats>,
) -> Self {
let reserved_bytes = reserved_bytes as u64;
stats.add_me_writer_byte_budget_queued_bytes(reserved_bytes);
Self {
_permit: permit,
reserved_bytes,
stats,
state: WriterBytePermitState::Queued,
}
}
/// Moves the reservation from queued to in-flight accounting exactly once.
pub(crate) fn mark_inflight(&mut self) {
if matches!(self.state, WriterBytePermitState::Queued) {
self.stats
.move_me_writer_byte_budget_to_inflight(self.reserved_bytes);
self.state = WriterBytePermitState::Inflight;
}
}
}
impl Drop for WriterBytePermit {
fn drop(&mut self) {
match self.state {
WriterBytePermitState::Queued => self
.stats
.release_me_writer_byte_budget_queued_bytes(self.reserved_bytes),
WriterBytePermitState::Inflight => self
.stats
.release_me_writer_byte_budget_inflight_bytes(self.reserved_bytes),
}
}
}
/// Commands sent to dedicated writer tasks to avoid mutex contention on TCP writes. /// Commands sent to dedicated writer tasks to avoid mutex contention on TCP writes.
pub(crate) enum WriterCommand { pub(crate) enum WriterCommand {
Data { Data {
payload: Bytes, payload: Bytes,
_permit: Option<OwnedSemaphorePermit>, _permit: Option<OwnedSemaphorePermit>,
writer_permit: WriterBytePermit,
}, },
DataAndFlush(Bytes), DataAndFlush(Bytes),
ProxyReq(ProxyReqCommand), ProxyReq(ProxyReqCommand),
@@ -33,6 +89,7 @@ pub(crate) struct ProxyReqCommand {
pub(crate) proxy_tag: Option<[u8; 16]>, pub(crate) proxy_tag: Option<[u8; 16]>,
pub(crate) payload: PooledBuffer, pub(crate) payload: PooledBuffer,
pub(crate) _permit: OwnedSemaphorePermit, pub(crate) _permit: OwnedSemaphorePermit,
pub(crate) writer_permit: WriterBytePermit,
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -87,7 +144,7 @@ fn build_rpc_frame_into(
) { ) {
let total_len = 4 + 4 + payload.len() + 4; let total_len = 4 + 4 + payload.len() + 4;
frame.clear(); frame.clear();
frame.reserve(total_len + 15); frame.reserve_exact(total_len + 15);
let total_len = total_len as u32; let total_len = total_len as u32;
frame.extend_from_slice(&total_len.to_le_bytes()); frame.extend_from_slice(&total_len.to_le_bytes());
frame.extend_from_slice(&seq_no.to_le_bytes()); frame.extend_from_slice(&seq_no.to_le_bytes());
@@ -279,7 +336,7 @@ impl RpcWriter {
); );
let total_len = 4 + 4 + payload_len + 4; let total_len = 4 + 4 + payload_len + 4;
self.frame_buf.clear(); self.frame_buf.clear();
self.frame_buf.reserve(total_len + 15); self.frame_buf.reserve_exact(total_len + 15);
self.frame_buf self.frame_buf
.extend_from_slice(&(total_len as u32).to_le_bytes()); .extend_from_slice(&(total_len as u32).to_le_bytes());
self.frame_buf.extend_from_slice(&self.seq_no.to_le_bytes()); self.frame_buf.extend_from_slice(&self.seq_no.to_le_bytes());
@@ -327,3 +384,35 @@ impl RpcWriter {
self.writer.flush().await.map_err(ProxyError::Io) self.writer.flush().await.map_err(ProxyError::Io)
} }
} }
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::Semaphore;
#[test]
fn writer_byte_permit_tracks_queued_and_inflight_lifecycle() {
let stats = Arc::new(Stats::default());
let semaphore = Arc::new(Semaphore::new(2));
let permit = semaphore
.clone()
.try_acquire_many_owned(2)
.expect("writer byte permits must be available");
let mut writer_permit = WriterBytePermit::new(permit, 32 * 1024, stats.clone());
assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 32 * 1024);
assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0);
writer_permit.mark_inflight();
assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0);
assert_eq!(
stats.get_me_writer_byte_budget_inflight_bytes_gauge(),
32 * 1024
);
drop(writer_permit);
assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0);
assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0);
assert_eq!(semaphore.available_permits(), 2);
}
}
+11 -2
View File
@@ -1795,6 +1795,7 @@ mod tests {
general.me_writer_pick_sample_size, general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(), MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity, general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity, general.me_route_channel_capacity,
general.me_route_backpressure_enabled, general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled, general.me_route_fairshare_enabled,
@@ -1821,6 +1822,7 @@ mod tests {
) -> u64 { ) -> u64 {
let (conn_id, _rx) = pool.registry.register().await; let (conn_id, _rx) = pool.registry.register().await;
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8); let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter { let writer = MeWriter {
id: writer_id, id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4000 + writer_id as u16), addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4000 + writer_id as u16),
@@ -1830,6 +1832,7 @@ mod tests {
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id), created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)), degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -1839,7 +1842,9 @@ mod tests {
allow_drain_fallback: Arc::new(AtomicBool::new(false)), allow_drain_fallback: Arc::new(AtomicBool::new(false)),
}; };
pool.writers.write().await.push(writer); pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await; pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed); pool.conn_count.fetch_add(1, Ordering::Relaxed);
assert!( assert!(
pool.registry pool.registry
@@ -1860,6 +1865,7 @@ mod tests {
async fn insert_live_writer(pool: &Arc<MePool>, writer_id: u64, writer_dc: i32) { async fn insert_live_writer(pool: &Arc<MePool>, writer_id: u64, writer_dc: i32) {
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8); let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter { let writer = MeWriter {
id: writer_id, id: writer_id,
addr: SocketAddr::new( addr: SocketAddr::new(
@@ -1877,6 +1883,7 @@ mod tests {
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
created_at: Instant::now(), created_at: Instant::now(),
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)), degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -1886,7 +1893,9 @@ mod tests {
allow_drain_fallback: Arc::new(AtomicBool::new(false)), allow_drain_fallback: Arc::new(AtomicBool::new(false)),
}; };
pool.writers.write().await.push(writer); pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await; pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed); pool.conn_count.fetch_add(1, Ordering::Relaxed);
} }
+16 -1
View File
@@ -10,7 +10,7 @@ use std::sync::atomic::{
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use arc_swap::ArcSwap; use arc_swap::ArcSwap;
use tokio::sync::{Mutex, RwLock, mpsc, watch}; use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch};
use tokio_util::sync::CancellationToken; use tokio_util::sync::CancellationToken;
use crate::config::{ use crate::config::{
@@ -48,6 +48,8 @@ pub struct MeWriter {
pub contour: Arc<AtomicU8>, pub contour: Arc<AtomicU8>,
pub created_at: Instant, pub created_at: Instant,
pub tx: mpsc::Sender<WriterCommand>, pub tx: mpsc::Sender<WriterCommand>,
/// Aggregate resident-memory budget shared by all data commands for this writer.
pub byte_budget: Arc<Semaphore>,
pub cancel: CancellationToken, pub cancel: CancellationToken,
pub degraded: Arc<AtomicBool>, pub degraded: Arc<AtomicBool>,
pub rtt_ema_ms_x10: Arc<AtomicU32>, pub rtt_ema_ms_x10: Arc<AtomicU32>,
@@ -277,6 +279,7 @@ pub(super) struct WriterLifecycleCore {
pub(super) me_keepalive_payload_random: bool, pub(super) me_keepalive_payload_random: bool,
pub(super) rpc_proxy_req_every_secs: AtomicU64, pub(super) rpc_proxy_req_every_secs: AtomicU64,
pub(super) writer_cmd_channel_capacity: usize, pub(super) writer_cmd_channel_capacity: usize,
pub(super) writer_byte_budget_permits: usize,
} }
pub(super) struct RouteRuntimeCore { pub(super) struct RouteRuntimeCore {
@@ -553,6 +556,7 @@ impl MePool {
me_writer_pick_sample_size: u8, me_writer_pick_sample_size: u8,
me_socks_kdf_policy: MeSocksKdfPolicy, me_socks_kdf_policy: MeSocksKdfPolicy,
me_writer_cmd_channel_capacity: usize, me_writer_cmd_channel_capacity: usize,
me_writer_byte_budget_bytes: usize,
me_route_channel_capacity: usize, me_route_channel_capacity: usize,
me_route_backpressure_enabled: bool, me_route_backpressure_enabled: bool,
me_route_fairshare_enabled: bool, me_route_fairshare_enabled: bool,
@@ -583,6 +587,7 @@ impl MePool {
); );
let (writer_epoch, _) = watch::channel(0u64); let (writer_epoch, _) = watch::channel(0u64);
let now_epoch_secs = Self::now_epoch_secs(); let now_epoch_secs = Self::now_epoch_secs();
stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes);
Arc::new(Self { Arc::new(Self {
routing: Arc::new(RoutingCore { routing: Arc::new(RoutingCore {
registry, registry,
@@ -615,6 +620,9 @@ impl MePool {
me_keepalive_payload_random, me_keepalive_payload_random,
rpc_proxy_req_every_secs: AtomicU64::new(rpc_proxy_req_every_secs), rpc_proxy_req_every_secs: AtomicU64::new(rpc_proxy_req_every_secs),
writer_cmd_channel_capacity: me_writer_cmd_channel_capacity.max(1), writer_cmd_channel_capacity: me_writer_cmd_channel_capacity.max(1),
writer_byte_budget_permits: me_writer_byte_budget_bytes
.div_ceil(crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
.max(1),
}), }),
route_runtime: Arc::new(RouteRuntimeCore { route_runtime: Arc::new(RouteRuntimeCore {
me_route_no_writer_mode: AtomicU8::new(me_route_no_writer_mode.as_u8()), me_route_no_writer_mode: AtomicU8::new(me_route_no_writer_mode.as_u8()),
@@ -833,6 +841,13 @@ impl MePool {
}) })
} }
/// Creates the immutable byte semaphore assigned to one ME writer generation.
pub(crate) fn new_writer_byte_budget(&self) -> Arc<Semaphore> {
Arc::new(Semaphore::new(
self.writer_lifecycle.writer_byte_budget_permits,
))
}
pub fn current_generation(&self) -> u64 { pub fn current_generation(&self) -> u64 {
self.reinit.active_generation.load(Ordering::Relaxed) self.reinit.active_generation.load(Ordering::Relaxed)
} }
+13 -3
View File
@@ -63,13 +63,19 @@ async fn writer_command_loop(
_ = cancel.cancelled() => return Ok(()), _ = cancel.cancelled() => return Ok(()),
cmd = rx.recv() => { cmd = rx.recv() => {
match cmd { match cmd {
Some(WriterCommand::Data { payload, _permit }) => { Some(WriterCommand::Data {
payload,
_permit,
mut writer_permit,
}) => {
writer_permit.mark_inflight();
rpc_writer.send(&payload).await?; rpc_writer.send(&payload).await?;
} }
Some(WriterCommand::DataAndFlush(payload)) => { Some(WriterCommand::DataAndFlush(payload)) => {
rpc_writer.send_and_flush(&payload).await?; rpc_writer.send_and_flush(&payload).await?;
} }
Some(WriterCommand::ProxyReq(command)) => { Some(WriterCommand::ProxyReq(mut command)) => {
command.writer_permit.mark_inflight();
rpc_writer.send_proxy_req(&command).await?; rpc_writer.send_proxy_req(&command).await?;
} }
Some(WriterCommand::ControlAndFlush(payload)) => { Some(WriterCommand::ControlAndFlush(payload)) => {
@@ -419,6 +425,7 @@ impl MePool {
let draining_started_at_epoch_secs = Arc::new(AtomicU64::new(0)); let draining_started_at_epoch_secs = Arc::new(AtomicU64::new(0));
let drain_deadline_epoch_secs = Arc::new(AtomicU64::new(0)); let drain_deadline_epoch_secs = Arc::new(AtomicU64::new(0));
let allow_drain_fallback = Arc::new(AtomicBool::new(false)); let allow_drain_fallback = Arc::new(AtomicBool::new(false));
let byte_budget = self.new_writer_byte_budget();
let (tx, rx) = let (tx, rx) =
mpsc::channel::<WriterCommand>(self.writer_lifecycle.writer_cmd_channel_capacity); mpsc::channel::<WriterCommand>(self.writer_lifecycle.writer_cmd_channel_capacity);
let rpc_writer = RpcWriter { let rpc_writer = RpcWriter {
@@ -438,6 +445,7 @@ impl MePool {
contour: contour.clone(), contour: contour.clone(),
created_at: Instant::now(), created_at: Instant::now(),
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: cancel.clone(), cancel: cancel.clone(),
degraded: degraded.clone(), degraded: degraded.clone(),
rtt_ema_ms_x10: rtt_ema_ms_x10.clone(), rtt_ema_ms_x10: rtt_ema_ms_x10.clone(),
@@ -449,7 +457,9 @@ impl MePool {
self.writers self.writers
.update(|writers| writers.push(writer.clone())) .update(|writers| writers.push(writer.clone()))
.await; .await;
self.registry.register_writer(writer_id, tx.clone()).await; self.registry
.register_writer(writer_id, tx.clone(), byte_budget)
.await;
self.registry.mark_writer_idle(writer_id).await; self.registry.mark_writer_idle(writer_id).await;
self.conn_count.fetch_add(1, Ordering::Relaxed); self.conn_count.fetch_add(1, Ordering::Relaxed);
self.notify_writer_epoch(); self.notify_writer_epoch();
+9 -1
View File
@@ -48,6 +48,8 @@ pub struct BoundConn {
pub struct ConnWriter { pub struct ConnWriter {
pub writer_id: u64, pub writer_id: u64,
pub tx: mpsc::Sender<WriterCommand>, pub tx: mpsc::Sender<WriterCommand>,
/// Writer-local memory budget used by the hot bound-client route.
pub byte_budget: Arc<Semaphore>,
} }
#[derive(Clone, Debug, Default)] #[derive(Clone, Debug, Default)]
@@ -62,7 +64,13 @@ struct RoutingTable {
} }
struct WriterTable { struct WriterTable {
map: DashMap<u64, mpsc::Sender<WriterCommand>>, map: DashMap<u64, WriterRoute>,
}
#[derive(Clone)]
struct WriterRoute {
tx: mpsc::Sender<WriterCommand>,
byte_budget: Arc<Semaphore>,
} }
#[derive(Clone)] #[derive(Clone)]
+30 -8
View File
@@ -1,10 +1,16 @@
use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;
use bytes::Bytes; use bytes::Bytes;
use tokio::sync::Semaphore;
use super::{ConnMeta, ConnRegistry, RouteResult}; use super::{ConnMeta, ConnRegistry, RouteResult};
use crate::transport::middle_proxy::MeResponse; use crate::transport::middle_proxy::MeResponse;
fn writer_byte_budget() -> Arc<Semaphore> {
Arc::new(Semaphore::new(2049))
}
#[tokio::test] #[tokio::test]
async fn writer_activity_snapshot_tracks_writer_and_dc_load() { async fn writer_activity_snapshot_tracks_writer_and_dc_load() {
let registry = ConnRegistry::new(); let registry = ConnRegistry::new();
@@ -14,8 +20,12 @@ async fn writer_activity_snapshot_tracks_writer_and_dc_load() {
let (conn_c, _rx_c) = registry.register().await; let (conn_c, _rx_c) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a.clone()).await; registry
registry.register_writer(20, writer_tx_b.clone()).await; .register_writer(10, writer_tx_a.clone(), writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b.clone(), writer_byte_budget())
.await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
assert!( assert!(
@@ -124,8 +134,12 @@ async fn bind_writer_rebinds_conn_atomically() {
let (conn_id, _rx) = registry.register().await; let (conn_id, _rx) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a).await; registry
registry.register_writer(20, writer_tx_b).await; .register_writer(10, writer_tx_a, writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b, writer_byte_budget())
.await;
let client_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); let client_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
let first_our_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443); let first_our_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443);
@@ -184,8 +198,12 @@ async fn writer_lost_does_not_drop_rebound_conn() {
let (conn_id, _rx) = registry.register().await; let (conn_id, _rx) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a).await; registry
registry.register_writer(20, writer_tx_b).await; .register_writer(10, writer_tx_a, writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b, writer_byte_budget())
.await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
assert!( assert!(
@@ -262,8 +280,12 @@ async fn non_empty_writer_ids_returns_only_writers_with_bound_clients() {
let (conn_id, _rx) = registry.register().await; let (conn_id, _rx) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a).await; registry
registry.register_writer(20, writer_tx_b).await; .register_writer(10, writer_tx_a, writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b, writer_byte_budget())
.await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
assert!( assert!(
+19 -4
View File
@@ -1,4 +1,5 @@
use std::collections::{HashMap, HashSet}; use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
use std::time::Duration; use std::time::Duration;
@@ -56,7 +57,13 @@ impl ConnRegistry {
} }
} }
pub async fn register_writer(&self, writer_id: u64, tx: mpsc::Sender<WriterCommand>) { /// Registers one writer command route and its matching memory budget atomically.
pub async fn register_writer(
&self,
writer_id: u64,
tx: mpsc::Sender<WriterCommand>,
byte_budget: Arc<tokio::sync::Semaphore>,
) {
let mut binding = self.binding.inner.lock().await; let mut binding = self.binding.inner.lock().await;
binding binding
.conns_for_writer .conns_for_writer
@@ -70,7 +77,13 @@ impl ConnRegistry {
.writer_idle_since_epoch_secs .writer_idle_since_epoch_secs
.entry(writer_id) .entry(writer_id)
.or_insert_with(Self::now_epoch_secs); .or_insert_with(Self::now_epoch_secs);
self.writers.map.insert(writer_id, tx); self.writers.map.insert(
writer_id,
super::WriterRoute {
tx,
byte_budget,
},
);
} }
/// Unregister connection, returning associated writer_id if any. /// Unregister connection, returning associated writer_id if any.
@@ -417,7 +430,8 @@ impl ConnRegistry {
.map(|entry| entry.value().clone())?; .map(|entry| entry.value().clone())?;
Some(ConnWriter { Some(ConnWriter {
writer_id, writer_id,
tx: writer, tx: writer.tx,
byte_budget: writer.byte_budget,
}) })
} }
@@ -438,7 +452,8 @@ impl ConnRegistry {
Some(( Some((
ConnWriter { ConnWriter {
writer_id, writer_id,
tx: writer, tx: writer.tx,
byte_budget: writer.byte_budget,
}, },
meta, meta,
)) ))
+299 -50
View File
@@ -6,16 +6,18 @@ use std::sync::Arc;
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use tokio::sync::{OwnedSemaphorePermit, mpsc}; use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc};
use tokio::sync::mpsc::error::TrySendError; use tokio::sync::mpsc::error::TrySendError;
use tracing::{debug, warn}; use tracing::{debug, warn};
use super::MePool; use super::MePool;
use super::codec::{ProxyReqCommand, WriterCommand}; use super::codec::{ProxyReqCommand, WriterBytePermit, WriterCommand};
use super::registry::ConnMeta; use super::registry::ConnMeta;
use super::wire::build_proxy_req_payload; use super::wire::{build_proxy_req_payload, proxy_req_payload_len};
use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
use crate::config::{MeRouteNoWriterMode, MeWriterPickMode}; use crate::config::{MeRouteNoWriterMode, MeWriterPickMode};
use crate::error::{ProxyError, Result}; use crate::error::{ProxyError, Result};
use crate::stats::Stats;
use crate::stream::PooledBuffer; use crate::stream::PooledBuffer;
use rand::seq::SliceRandom; use rand::seq::SliceRandom;
@@ -29,6 +31,8 @@ const PICK_PENALTY_WARM: u64 = 200;
const PICK_PENALTY_DRAINING: u64 = 600; const PICK_PENALTY_DRAINING: u64 = 600;
const PICK_PENALTY_STALE: u64 = 300; const PICK_PENALTY_STALE: u64 = 300;
const PICK_PENALTY_DEGRADED: u64 = 250; const PICK_PENALTY_DEGRADED: u64 = 250;
const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27;
const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128;
mod close; mod close;
mod recovery; mod recovery;
@@ -39,6 +43,11 @@ enum WriterCommandReserveError {
TimedOut, TimedOut,
} }
enum WriterByteReserveError {
Closed,
TimedOut,
}
fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> { fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> {
tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok()) tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok())
} }
@@ -61,11 +70,16 @@ fn payload_permit_from_data_command(cmd: WriterCommand) -> Option<OwnedSemaphore
async fn reserve_writer_command_slot( async fn reserve_writer_command_slot(
tx: &mpsc::Sender<WriterCommand>, tx: &mpsc::Sender<WriterCommand>,
wait: Option<Duration>, deadline: Option<Instant>,
) -> std::result::Result<mpsc::OwnedPermit<WriterCommand>, WriterCommandReserveError> { ) -> std::result::Result<mpsc::OwnedPermit<WriterCommand>, WriterCommandReserveError> {
let reserve = tx.clone().reserve_owned(); let reserve = tx.clone().reserve_owned();
match wait { match deadline {
Some(wait) => match tokio::time::timeout(wait, reserve).await { Some(deadline) => match tokio::time::timeout(
deadline.saturating_duration_since(Instant::now()),
reserve,
)
.await
{
Ok(Ok(permit)) => Ok(permit), Ok(Ok(permit)) => Ok(permit),
Ok(Err(_)) => Err(WriterCommandReserveError::Closed), Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
Err(_) => Err(WriterCommandReserveError::TimedOut), Err(_) => Err(WriterCommandReserveError::TimedOut),
@@ -74,6 +88,88 @@ async fn reserve_writer_command_slot(
} }
} }
fn writer_send_deadline(wait: Option<Duration>) -> Option<Instant> {
wait.map(|wait| Instant::now() + wait)
}
fn writer_resident_permits(
source_capacity: usize,
encoded_payload_len: usize,
) -> Option<(u32, usize)> {
let resident_bytes = source_capacity
.checked_add(encoded_payload_len)?
.checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?;
let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES);
let permits = u32::try_from(permits).ok()?;
let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?;
Some((permits.max(1), reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)))
}
fn proxy_req_resident_permits(
source_capacity: usize,
data_len: usize,
proxy_tag: Option<&[u8]>,
proto_flags: u32,
) -> Option<(u32, usize)> {
writer_resident_permits(
source_capacity,
proxy_req_payload_len(data_len, proxy_tag, proto_flags),
)
}
fn try_reserve_writer_bytes(
byte_budget: &Arc<Semaphore>,
permits: u32,
reserved_bytes: usize,
stats: &Arc<Stats>,
) -> std::result::Result<WriterBytePermit, TryAcquireError> {
byte_budget
.clone()
.try_acquire_many_owned(permits)
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
}
async fn reserve_writer_bytes(
byte_budget: &Arc<Semaphore>,
permits: u32,
reserved_bytes: usize,
deadline: Option<Instant>,
stats: &Arc<Stats>,
) -> std::result::Result<WriterBytePermit, WriterByteReserveError> {
match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) {
Ok(permit) => return Ok(permit),
Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed),
Err(TryAcquireError::NoPermits) => {
stats.increment_me_writer_byte_budget_wait_total();
}
}
let acquire = byte_budget.clone().acquire_many_owned(permits);
match deadline {
Some(deadline) => match tokio::time::timeout(
deadline.saturating_duration_since(Instant::now()),
acquire,
)
.await
{
Ok(Ok(permit)) => Ok(WriterBytePermit::new(
permit,
reserved_bytes,
stats.clone(),
)),
Ok(Err(_)) => Err(WriterByteReserveError::Closed),
Err(_) => {
stats.increment_me_writer_byte_budget_timeout_total();
Err(WriterByteReserveError::TimedOut)
}
},
None => acquire
.await
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
.map_err(|_| WriterByteReserveError::Closed),
}
}
impl MePool { impl MePool {
/// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default. /// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default.
/// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command. /// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command.
@@ -89,6 +185,35 @@ impl MePool {
mut payload_permit: Option<OwnedSemaphorePermit>, mut payload_permit: Option<OwnedSemaphorePermit>,
) -> Result<()> { ) -> Result<()> {
let tag = tag_override.or(self.proxy_tag.as_deref()); let tag = tag_override.or(self.proxy_tag.as_deref());
let Some(source_capacity) = data
.len()
.checked_add(LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES)
else {
self.stats
.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload residency calculation overflow".into(),
));
};
let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits(
source_capacity,
data.len(),
tag,
proto_flags,
) else {
self.stats
.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload residency calculation overflow".into(),
));
};
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
self.stats
.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload exceeds configured byte budget".into(),
));
}
let build_routed_payload = |effective_our_addr: SocketAddr| { let build_routed_payload = |effective_our_addr: SocketAddr| {
( (
build_proxy_req_payload( build_proxy_req_payload(
@@ -129,10 +254,38 @@ impl MePool {
loop { loop {
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await
{ {
let deadline = writer_send_deadline(
self.route_runtime.me_route_blocking_send_timeout,
);
let writer_permit = match reserve_writer_bytes(
&current.byte_budget,
writer_byte_permits,
writer_reserved_bytes,
deadline,
&self.stats,
)
.await
{
Ok(permit) => permit,
Err(WriterByteReserveError::TimedOut) => {
self.stats
.increment_me_writer_pick_full_total(self.writer_pick_mode());
return Err(ProxyError::Proxy(
"ME writer byte budget full within blocking send timeout".into(),
));
}
Err(WriterByteReserveError::Closed) => {
warn!(writer_id = current.writer_id, "ME writer byte budget closed");
self.remove_writer_and_close_clients(current.writer_id)
.await;
continue;
}
};
let (current_payload, _) = build_routed_payload(current_meta.our_addr); let (current_payload, _) = build_routed_payload(current_meta.our_addr);
let command = WriterCommand::Data { let command = WriterCommand::Data {
payload: current_payload, payload: current_payload,
_permit: payload_permit.take(), _permit: payload_permit.take(),
writer_permit,
}; };
match current.tx.try_send(command) { match current.tx.try_send(command) {
Ok(()) => { Ok(()) => {
@@ -142,7 +295,7 @@ impl MePool {
Err(TrySendError::Full(cmd)) => { Err(TrySendError::Full(cmd)) => {
match reserve_writer_command_slot( match reserve_writer_command_slot(
&current.tx, &current.tx,
self.route_runtime.me_route_blocking_send_timeout, deadline,
) )
.await .await
{ {
@@ -482,11 +635,31 @@ impl MePool {
if !self.writer_accepts_new_binding(w) { if !self.writer_accepts_new_binding(w) {
continue; continue;
} }
// Keep the advertised proxy IP aligned with the selected ME writer source. let writer_permit = match try_reserve_writer_bytes(
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); &w.byte_budget,
let (payload, meta) = build_routed_payload(effective_our_addr); writer_byte_permits,
writer_reserved_bytes,
&self.stats,
) {
Ok(permit) => permit,
Err(TryAcquireError::NoPermits) => {
if fallback_blocking_idx.is_none() {
fallback_blocking_idx = Some(idx);
}
continue;
}
Err(TryAcquireError::Closed) => {
self.stats.increment_me_writer_pick_closed_total(pick_mode);
warn!(writer_id = w.id, "ME writer byte budget closed");
self.remove_writer_and_close_clients(w.id).await;
continue;
}
};
match w.tx.clone().try_reserve_owned() { match w.tx.clone().try_reserve_owned() {
Ok(permit) => { Ok(permit) => {
// Keep the advertised proxy IP aligned with the selected ME writer source.
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
let (payload, meta) = build_routed_payload(effective_our_addr);
if !self.registry.bind_writer(conn_id, w.id, meta).await { if !self.registry.bind_writer(conn_id, w.id, meta).await {
debug!( debug!(
conn_id, conn_id,
@@ -500,6 +673,7 @@ impl MePool {
permit.send(WriterCommand::Data { permit.send(WriterCommand::Data {
payload, payload,
_permit: payload_permit.take(), _permit: payload_permit.take(),
writer_permit,
}); });
self.stats self.stats
.increment_me_writer_pick_success_try_total(pick_mode); .increment_me_writer_pick_success_try_total(pick_mode);
@@ -542,51 +716,68 @@ impl MePool {
} }
self.stats self.stats
.increment_me_writer_pick_blocking_fallback_total(); .increment_me_writer_pick_blocking_fallback_total();
// Keep the advertised proxy IP aligned with the selected ME writer source. let deadline = writer_send_deadline(
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); self.route_runtime.me_route_blocking_send_timeout,
let (payload, meta) = build_routed_payload(effective_our_addr); );
let reserve_result = let writer_permit = match reserve_writer_bytes(
if let Some(timeout) = self.route_runtime.me_route_blocking_send_timeout { &w.byte_budget,
match tokio::time::timeout(timeout, w.tx.clone().reserve_owned()).await { writer_byte_permits,
Ok(result) => result, writer_reserved_bytes,
Err(_) => { deadline,
self.stats.increment_me_writer_pick_full_total(pick_mode); &self.stats,
continue; )
} .await
} {
} else { Ok(permit) => permit,
w.tx.clone().reserve_owned().await Err(WriterByteReserveError::TimedOut) => {
}; self.stats.increment_me_writer_pick_full_total(pick_mode);
match reserve_result { continue;
Ok(permit) => {
if !self.registry.bind_writer(conn_id, w.id, meta).await {
debug!(
conn_id,
writer_id = w.id,
"ME writer disappeared before fallback bind commit, pruning stale writer"
);
drop(permit);
self.remove_writer_and_close_clients(w.id).await;
continue;
}
permit.send(WriterCommand::Data {
payload,
_permit: payload_permit.take(),
});
self.stats
.increment_me_writer_pick_success_fallback_total(pick_mode);
if w.generation < self.current_generation() {
self.stats.increment_pool_stale_pick_total();
}
self.note_hybrid_route_success();
return Ok(());
} }
Err(_) => { Err(WriterByteReserveError::Closed) => {
self.stats.increment_me_writer_pick_closed_total(pick_mode);
warn!(writer_id = w.id, "ME writer byte budget closed (blocking)");
self.remove_writer_and_close_clients(w.id).await;
continue;
}
};
let permit = match reserve_writer_command_slot(&w.tx, deadline).await {
Ok(permit) => permit,
Err(WriterCommandReserveError::TimedOut) => {
self.stats.increment_me_writer_pick_full_total(pick_mode);
continue;
}
Err(WriterCommandReserveError::Closed) => {
self.stats.increment_me_writer_pick_closed_total(pick_mode); self.stats.increment_me_writer_pick_closed_total(pick_mode);
warn!(writer_id = w.id, "ME writer channel closed (blocking)"); warn!(writer_id = w.id, "ME writer channel closed (blocking)");
self.remove_writer_and_close_clients(w.id).await; self.remove_writer_and_close_clients(w.id).await;
continue;
} }
};
// Keep the advertised proxy IP aligned with the selected ME writer source.
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
let (payload, meta) = build_routed_payload(effective_our_addr);
if !self.registry.bind_writer(conn_id, w.id, meta).await {
debug!(
conn_id,
writer_id = w.id,
"ME writer disappeared before fallback bind commit, pruning stale writer"
);
drop(permit);
self.remove_writer_and_close_clients(w.id).await;
continue;
} }
permit.send(WriterCommand::Data {
payload,
_permit: payload_permit.take(),
writer_permit,
});
self.stats
.increment_me_writer_pick_success_fallback_total(pick_mode);
if w.generation < self.current_generation() {
self.stats.increment_pool_stale_pick_total();
}
self.note_hybrid_route_success();
return Ok(());
} }
} }
@@ -604,8 +795,65 @@ impl MePool {
tag_override: Option<[u8; 16]>, tag_override: Option<[u8; 16]>,
) -> Result<()> { ) -> Result<()> {
let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref())); let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref()));
let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits(
payload.capacity(),
payload.len(),
tag.as_ref().map(|tag| tag.as_slice()),
proto_flags,
) else {
self.stats
.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload residency calculation overflow".into(),
));
};
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
self.stats
.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload exceeds configured byte budget".into(),
));
}
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await {
let deadline = writer_send_deadline(
self.route_runtime.me_route_blocking_send_timeout,
);
let writer_permit = match reserve_writer_bytes(
&current.byte_budget,
writer_byte_permits,
writer_reserved_bytes,
deadline,
&self.stats,
)
.await
{
Ok(permit) => permit,
Err(WriterByteReserveError::TimedOut) => {
self.stats
.increment_me_writer_pick_full_total(self.writer_pick_mode());
return Err(ProxyError::Proxy(
"ME writer byte budget full within blocking send timeout".into(),
));
}
Err(WriterByteReserveError::Closed) => {
warn!(writer_id = current.writer_id, "ME writer byte budget closed");
self.remove_writer_and_close_clients(current.writer_id)
.await;
return self
.send_proxy_req(
conn_id,
target_dc,
client_addr,
our_addr,
payload.as_ref(),
proto_flags,
tag.as_ref().map(|tag| tag.as_slice()),
Some(_permit),
)
.await;
}
};
let command = WriterCommand::ProxyReq(ProxyReqCommand { let command = WriterCommand::ProxyReq(ProxyReqCommand {
conn_id, conn_id,
client_addr, client_addr,
@@ -614,6 +862,7 @@ impl MePool {
proxy_tag: tag, proxy_tag: tag,
payload, payload,
_permit, _permit,
writer_permit,
}); });
match current.tx.try_send(command) { match current.tx.try_send(command) {
Ok(()) => { Ok(()) => {
@@ -623,7 +872,7 @@ impl MePool {
Err(TrySendError::Full(cmd)) => { Err(TrySendError::Full(cmd)) => {
match reserve_writer_command_slot( match reserve_writer_command_slot(
&current.tx, &current.tx,
self.route_runtime.me_route_blocking_send_timeout, deadline,
) )
.await .await
{ {
+12 -6
View File
@@ -10,7 +10,7 @@ use crate::protocol::constants::{RPC_CLOSE_CONN_U32, RPC_CLOSE_EXT_U32};
use super::super::MePool; use super::super::MePool;
use super::super::codec::{WriterCommand, build_control_payload}; use super::super::codec::{WriterCommand, build_control_payload};
use super::{WriterCommandReserveError, reserve_writer_command_slot}; use super::{WriterCommandReserveError, reserve_writer_command_slot, writer_send_deadline};
const ME_CLOSE_SIGNAL_SEND_TIMEOUT: Duration = Duration::from_millis(50); const ME_CLOSE_SIGNAL_SEND_TIMEOUT: Duration = Duration::from_millis(50);
@@ -22,8 +22,11 @@ impl MePool {
match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) { match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) {
Ok(()) => {} Ok(()) => {}
Err(TrySendError::Full(cmd)) => { Err(TrySendError::Full(cmd)) => {
match reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)) match reserve_writer_command_slot(
.await &w.tx,
writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)),
)
.await
{ {
Ok(permit) => { Ok(permit) => {
permit.send(cmd); permit.send(cmd);
@@ -63,9 +66,12 @@ impl MePool {
match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) { match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) {
Ok(()) => {} Ok(()) => {}
Err(TrySendError::Full(cmd)) => { Err(TrySendError::Full(cmd)) => {
let _ = reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)) let _ = reserve_writer_command_slot(
.await &w.tx,
.map(|permit| permit.send(cmd)); writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)),
)
.await
.map(|permit| permit.send(cmd));
} }
Err(TrySendError::Closed(_)) => { Err(TrySendError::Closed(_)) => {
debug!(conn_id, "ME close_conn skipped: writer channel closed"); debug!(conn_id, "ME close_conn skipped: writer channel closed");
@@ -105,6 +105,7 @@ async fn make_pool(
general.me_writer_pick_sample_size, general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(), MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity, general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity, general.me_route_channel_capacity,
general.me_route_backpressure_enabled, general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled, general.me_route_fairshare_enabled,
@@ -134,6 +135,7 @@ async fn insert_draining_writer(
drain_deadline_epoch_secs: u64, drain_deadline_epoch_secs: u64,
) { ) {
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8); let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter { let writer = MeWriter {
id: writer_id, id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 6000 + writer_id as u16), addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 6000 + writer_id as u16),
@@ -143,6 +145,7 @@ async fn insert_draining_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id), created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)), degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -153,7 +156,9 @@ async fn insert_draining_writer(
}; };
pool.writers.write().await.push(writer); pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await; pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed); pool.conn_count.fetch_add(1, Ordering::Relaxed);
for idx in 0..bound_clients { for idx in 0..bound_clients {
@@ -103,6 +103,7 @@ async fn make_pool(
general.me_writer_pick_sample_size, general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(), MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity, general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity, general.me_route_channel_capacity,
general.me_route_backpressure_enabled, general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled, general.me_route_fairshare_enabled,
@@ -131,6 +132,7 @@ async fn insert_draining_writer(
drain_deadline_epoch_secs: u64, drain_deadline_epoch_secs: u64,
) { ) {
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8); let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter { let writer = MeWriter {
id: writer_id, id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5500 + writer_id as u16), addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5500 + writer_id as u16),
@@ -140,6 +142,7 @@ async fn insert_draining_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id), created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)), degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -149,7 +152,9 @@ async fn insert_draining_writer(
allow_drain_fallback: Arc::new(AtomicBool::new(false)), allow_drain_fallback: Arc::new(AtomicBool::new(false)),
}; };
pool.writers.write().await.push(writer); pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await; pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed); pool.conn_count.fetch_add(1, Ordering::Relaxed);
for idx in 0..bound_clients { for idx in 0..bound_clients {
let (conn_id, _rx) = pool.registry.register().await; let (conn_id, _rx) = pool.registry.register().await;
@@ -98,6 +98,7 @@ async fn make_pool(me_pool_drain_threshold: u64) -> Arc<MePool> {
general.me_writer_pick_sample_size, general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(), MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity, general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity, general.me_route_channel_capacity,
general.me_route_backpressure_enabled, general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled, general.me_route_fairshare_enabled,
@@ -126,6 +127,7 @@ async fn insert_draining_writer(
) -> Vec<u64> { ) -> Vec<u64> {
let mut conn_ids = Vec::with_capacity(bound_clients); let mut conn_ids = Vec::with_capacity(bound_clients);
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8); let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter { let writer = MeWriter {
id: writer_id, id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4500 + writer_id as u16), addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4500 + writer_id as u16),
@@ -135,6 +137,7 @@ async fn insert_draining_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id), created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)), degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -144,7 +147,9 @@ async fn insert_draining_writer(
allow_drain_fallback: Arc::new(AtomicBool::new(false)), allow_drain_fallback: Arc::new(AtomicBool::new(false)),
}; };
pool.writers.write().await.push(writer); pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await; pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed); pool.conn_count.fetch_add(1, Ordering::Relaxed);
for idx in 0..bound_clients { for idx in 0..bound_clients {
let (conn_id, _rx) = pool.registry.register().await; let (conn_id, _rx) = pool.registry.register().await;
@@ -87,6 +87,7 @@ async fn make_pool() -> Arc<MePool> {
general.me_writer_pick_sample_size, general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(), MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity, general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity, general.me_route_channel_capacity,
general.me_route_backpressure_enabled, general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled, general.me_route_fairshare_enabled,
@@ -92,6 +92,7 @@ async fn make_pool() -> Arc<MePool> {
general.me_writer_pick_sample_size, general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(), MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity, general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity, general.me_route_channel_capacity,
general.me_route_backpressure_enabled, general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled, general.me_route_fairshare_enabled,
@@ -120,6 +121,7 @@ async fn insert_writer(
created_at: Instant, created_at: Instant,
) { ) {
let (tx, _rx) = mpsc::channel::<WriterCommand>(8); let (tx, _rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let contour = if draining { let contour = if draining {
WriterContour::Draining WriterContour::Draining
} else { } else {
@@ -134,6 +136,7 @@ async fn insert_writer(
contour: Arc::new(AtomicU8::new(contour.as_u8())), contour: Arc::new(AtomicU8::new(contour.as_u8())),
created_at, created_at,
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)), degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -144,7 +147,9 @@ async fn insert_writer(
}; };
pool.writers.write().await.push(writer); pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await; pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed); pool.conn_count.fetch_add(1, Ordering::Relaxed);
} }
@@ -98,6 +98,7 @@ async fn make_pool() -> (Arc<MePool>, Arc<SecureRandom>) {
general.me_writer_pick_sample_size, general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(), MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity, general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity, general.me_route_channel_capacity,
general.me_route_backpressure_enabled, general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled, general.me_route_fairshare_enabled,
@@ -127,6 +128,7 @@ async fn insert_writer(
register_in_registry: bool, register_in_registry: bool,
) -> mpsc::Receiver<WriterCommand> { ) -> mpsc::Receiver<WriterCommand> {
let (tx, rx) = mpsc::channel::<WriterCommand>(8); let (tx, rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter { let writer = MeWriter {
id: writer_id, id: writer_id,
addr, addr,
@@ -136,6 +138,7 @@ async fn insert_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
created_at: Instant::now(), created_at: Instant::now(),
tx: tx.clone(), tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(), cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)), degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -154,7 +157,9 @@ async fn insert_writer(
} }
pool.rebuild_endpoint_dc_map().await; pool.rebuild_endpoint_dc_map().await;
if register_in_registry { if register_in_registry {
pool.registry.register_writer(writer_id, tx).await; pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
} }
rx rx
} }