mirror of
https://github.com/telemt/telemt.git
synced 2026-07-15 16:30:20 +03:00
Merge pull request #875 from telemt/flow-aescr
Bound ME writer queues by resident payload bytes + Enable expanded AES key schedule zeroization
This commit is contained in:
Generated
+2
-1
@@ -21,6 +21,7 @@ dependencies = [
|
|||||||
"cfg-if",
|
"cfg-if",
|
||||||
"cipher",
|
"cipher",
|
||||||
"cpufeatures 0.2.17",
|
"cpufeatures 0.2.17",
|
||||||
|
"zeroize",
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
@@ -2899,7 +2900,7 @@ checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "telemt"
|
name = "telemt"
|
||||||
version = "3.4.23"
|
version = "3.4.24"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"aes",
|
"aes",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
|
|||||||
+3
-3
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "telemt"
|
name = "telemt"
|
||||||
version = "3.4.23"
|
version = "3.4.24"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
@@ -15,8 +15,8 @@ tokio = { version = "1.52.3", features = ["full", "tracing"] }
|
|||||||
tokio-util = { version = "0.7.18", features = ["full"] }
|
tokio-util = { version = "0.7.18", features = ["full"] }
|
||||||
|
|
||||||
# Crypto
|
# Crypto
|
||||||
aes = "0.8.4"
|
aes = { version = "0.8.4", features = ["zeroize"] }
|
||||||
ctr = "0.9.2"
|
ctr = { version = "0.9.2", features = ["zeroize"] }
|
||||||
cbc = "0.1.2"
|
cbc = "0.1.2"
|
||||||
sha2 = "0.10.9"
|
sha2 = "0.10.9"
|
||||||
sha1 = "0.10.6"
|
sha1 = "0.10.6"
|
||||||
|
|||||||
@@ -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;
|
||||||
@@ -35,6 +39,8 @@ const DEFAULT_ME_QUOTA_SOFT_OVERSHOOT_BYTES: u64 = 64 * 1024;
|
|||||||
const DEFAULT_ME_D2C_FRAME_BUF_SHRINK_THRESHOLD_BYTES: usize = 256 * 1024;
|
const DEFAULT_ME_D2C_FRAME_BUF_SHRINK_THRESHOLD_BYTES: usize = 256 * 1024;
|
||||||
const DEFAULT_DIRECT_RELAY_COPY_BUF_C2S_BYTES: usize = 64 * 1024;
|
const DEFAULT_DIRECT_RELAY_COPY_BUF_C2S_BYTES: usize = 64 * 1024;
|
||||||
const DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES: usize = 256 * 1024;
|
const DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES: usize = 256 * 1024;
|
||||||
|
pub(crate) const DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES: usize = 4 * 1024;
|
||||||
|
const DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES: usize = 0;
|
||||||
const DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE: u8 = 3;
|
const DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE: u8 = 3;
|
||||||
const DEFAULT_ME_HEALTH_INTERVAL_MS_UNHEALTHY: u64 = 1000;
|
const DEFAULT_ME_HEALTH_INTERVAL_MS_UNHEALTHY: u64 = 1000;
|
||||||
const DEFAULT_ME_HEALTH_INTERVAL_MS_HEALTHY: u64 = 3000;
|
const DEFAULT_ME_HEALTH_INTERVAL_MS_HEALTHY: u64 = 3000;
|
||||||
@@ -455,6 +461,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
|
||||||
}
|
}
|
||||||
@@ -499,6 +517,10 @@ pub(crate) fn default_direct_relay_copy_buf_s2c_bytes() -> usize {
|
|||||||
DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES
|
DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn default_direct_relay_buffer_budget_max_bytes() -> usize {
|
||||||
|
DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn default_me_writer_pick_sample_size() -> u8 {
|
pub(crate) fn default_me_writer_pick_sample_size() -> u8 {
|
||||||
DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE
|
DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -39,8 +39,11 @@ 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_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 16 * 1024 * 1024;
|
||||||
|
const MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 2 * 1024 * 1024 * 1024;
|
||||||
const MIN_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024;
|
const MIN_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024;
|
||||||
const MAX_MAX_CLIENT_FRAME_BYTES: usize = 16 * 1024 * 1024;
|
const MAX_MAX_CLIENT_FRAME_BYTES: usize = 16 * 1024 * 1024;
|
||||||
const MAX_API_REQUEST_BODY_LIMIT_BYTES: usize = 1024 * 1024;
|
const MAX_API_REQUEST_BODY_LIMIT_BYTES: usize = 1024 * 1024;
|
||||||
@@ -540,6 +543,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(),
|
||||||
@@ -599,6 +618,24 @@ impl ProxyConfig {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if config.general.direct_relay_buffer_budget_max_bytes != 0 {
|
||||||
|
if config.general.direct_relay_buffer_budget_max_bytes
|
||||||
|
% DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES
|
||||||
|
!= 0
|
||||||
|
{
|
||||||
|
return Err(ProxyError::Config(format!(
|
||||||
|
"general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of {DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
if !(MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES..=MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES)
|
||||||
|
.contains(&config.general.direct_relay_buffer_budget_max_bytes)
|
||||||
|
{
|
||||||
|
return Err(ProxyError::Config(format!(
|
||||||
|
"general.direct_relay_buffer_budget_max_bytes must be 0 or within [{MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES}, {MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES}]"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
if config.general.me_health_interval_ms_unhealthy == 0 {
|
if config.general.me_health_interval_ms_unhealthy == 0 {
|
||||||
return Err(ProxyError::Config(
|
return Err(ProxyError::Config(
|
||||||
"general.me_health_interval_ms_unhealthy must be > 0".to_string(),
|
"general.me_health_interval_ms_unhealthy must be > 0".to_string(),
|
||||||
|
|||||||
@@ -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",
|
||||||
@@ -64,6 +65,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
|
|||||||
"me_d2c_frame_buf_shrink_threshold_bytes",
|
"me_d2c_frame_buf_shrink_threshold_bytes",
|
||||||
"direct_relay_copy_buf_c2s_bytes",
|
"direct_relay_copy_buf_c2s_bytes",
|
||||||
"direct_relay_copy_buf_s2c_bytes",
|
"direct_relay_copy_buf_s2c_bytes",
|
||||||
|
"direct_relay_buffer_budget_max_bytes",
|
||||||
"crypto_pending_buffer",
|
"crypto_pending_buffer",
|
||||||
"max_client_frame",
|
"max_client_frame",
|
||||||
"desync_all_full",
|
"desync_all_full",
|
||||||
|
|||||||
@@ -95,6 +95,101 @@ 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]
|
||||||
|
fn load_rejects_unaligned_direct_relay_buffer_budget() {
|
||||||
|
let path = write_temp_config(
|
||||||
|
r#"
|
||||||
|
[general]
|
||||||
|
direct_relay_buffer_budget_max_bytes = 16777217
|
||||||
|
"#,
|
||||||
|
);
|
||||||
|
|
||||||
|
let err = ProxyConfig::load(&path).expect_err("unaligned direct relay buffer budget must fail");
|
||||||
|
assert!(
|
||||||
|
err.to_string().contains(
|
||||||
|
"general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of 4096"
|
||||||
|
)
|
||||||
|
);
|
||||||
|
remove_temp_config(&path);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn load_rejects_direct_relay_buffer_budget_above_hard_cap() {
|
||||||
|
let path = write_temp_config(
|
||||||
|
r#"
|
||||||
|
[general]
|
||||||
|
direct_relay_buffer_budget_max_bytes = 2147487744
|
||||||
|
"#,
|
||||||
|
);
|
||||||
|
|
||||||
|
let err =
|
||||||
|
ProxyConfig::load(&path).expect_err("direct relay buffer budget above hard cap must fail");
|
||||||
|
assert!(err.to_string().contains(
|
||||||
|
"general.direct_relay_buffer_budget_max_bytes must be 0 or within [16777216, 2147483648]"
|
||||||
|
));
|
||||||
|
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,16 +234,23 @@ 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
|
||||||
|
direct_relay_buffer_budget_max_bytes = 2147483648
|
||||||
max_client_frame = 16777216
|
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.direct_relay_buffer_budget_max_bytes,
|
||||||
|
2 * 1024 * 1024 * 1024
|
||||||
|
);
|
||||||
assert_eq!(cfg.general.max_client_frame, 16 * 1024 * 1024);
|
assert_eq!(cfg.general.max_client_frame, 16 * 1024 * 1024);
|
||||||
|
|
||||||
remove_temp_config(&path);
|
remove_temp_config(&path);
|
||||||
|
|||||||
+13
-2
@@ -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,
|
||||||
@@ -622,7 +626,7 @@ pub struct GeneralConfig {
|
|||||||
#[serde(default = "default_me_d2c_frame_buf_shrink_threshold_bytes")]
|
#[serde(default = "default_me_d2c_frame_buf_shrink_threshold_bytes")]
|
||||||
pub me_d2c_frame_buf_shrink_threshold_bytes: usize,
|
pub me_d2c_frame_buf_shrink_threshold_bytes: usize,
|
||||||
|
|
||||||
/// Copy buffer size for client->DC direction in direct relay.
|
/// Copy buffer ceiling for client->DC direction in direct relay.
|
||||||
///
|
///
|
||||||
/// This is also the upper bound for one amortized upload rate-limit burst:
|
/// This is also the upper bound for one amortized upload rate-limit burst:
|
||||||
/// upload debt is settled before the next relay read instead of blocking
|
/// upload debt is settled before the next relay read instead of blocking
|
||||||
@@ -630,13 +634,18 @@ pub struct GeneralConfig {
|
|||||||
#[serde(default = "default_direct_relay_copy_buf_c2s_bytes")]
|
#[serde(default = "default_direct_relay_copy_buf_c2s_bytes")]
|
||||||
pub direct_relay_copy_buf_c2s_bytes: usize,
|
pub direct_relay_copy_buf_c2s_bytes: usize,
|
||||||
|
|
||||||
/// Copy buffer size for DC->client direction in direct relay.
|
/// Copy buffer ceiling for DC->client direction in direct relay.
|
||||||
///
|
///
|
||||||
/// This bounds one direct download rate-limit grant because writes are
|
/// This bounds one direct download rate-limit grant because writes are
|
||||||
/// clipped to the currently available shaper budget.
|
/// clipped to the currently available shaper budget.
|
||||||
#[serde(default = "default_direct_relay_copy_buf_s2c_bytes")]
|
#[serde(default = "default_direct_relay_copy_buf_s2c_bytes")]
|
||||||
pub direct_relay_copy_buf_s2c_bytes: usize,
|
pub direct_relay_copy_buf_s2c_bytes: usize,
|
||||||
|
|
||||||
|
/// Process-wide hard ceiling for Direct relay copy buffers.
|
||||||
|
/// `0` derives the ceiling from host and cgroup memory limits.
|
||||||
|
#[serde(default = "default_direct_relay_buffer_budget_max_bytes")]
|
||||||
|
pub direct_relay_buffer_budget_max_bytes: usize,
|
||||||
|
|
||||||
/// Max pending ciphertext buffer per client writer (bytes).
|
/// Max pending ciphertext buffer per client writer (bytes).
|
||||||
/// Controls FakeTLS backpressure vs throughput.
|
/// Controls FakeTLS backpressure vs throughput.
|
||||||
#[serde(default = "default_crypto_pending_buffer")]
|
#[serde(default = "default_crypto_pending_buffer")]
|
||||||
@@ -1103,6 +1112,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(),
|
||||||
@@ -1116,6 +1126,7 @@ impl Default for GeneralConfig {
|
|||||||
default_me_d2c_frame_buf_shrink_threshold_bytes(),
|
default_me_d2c_frame_buf_shrink_threshold_bytes(),
|
||||||
direct_relay_copy_buf_c2s_bytes: default_direct_relay_copy_buf_c2s_bytes(),
|
direct_relay_copy_buf_c2s_bytes: default_direct_relay_copy_buf_c2s_bytes(),
|
||||||
direct_relay_copy_buf_s2c_bytes: default_direct_relay_copy_buf_s2c_bytes(),
|
direct_relay_copy_buf_s2c_bytes: default_direct_relay_copy_buf_s2c_bytes(),
|
||||||
|
direct_relay_buffer_budget_max_bytes: default_direct_relay_buffer_budget_max_bytes(),
|
||||||
me_warmup_stagger_enabled: default_true(),
|
me_warmup_stagger_enabled: default_true(),
|
||||||
me_warmup_step_delay_ms: default_warmup_step_delay_ms(),
|
me_warmup_step_delay_ms: default_warmup_step_delay_ms(),
|
||||||
me_warmup_step_jitter_ms: default_warmup_step_jitter_ms(),
|
me_warmup_step_jitter_ms: default_warmup_step_jitter_ms(),
|
||||||
|
|||||||
+18
-12
@@ -4,12 +4,12 @@
|
|||||||
//!
|
//!
|
||||||
//! ## Zeroize policy
|
//! ## Zeroize policy
|
||||||
//!
|
//!
|
||||||
//! - `AesCbc` stores raw key/IV bytes and zeroizes them on drop.
|
//! - `AesCbc` stores raw key/IV bytes and zeroizes them on drop. Temporary
|
||||||
//! - `AesCtr` wraps an opaque `Aes256Ctr` cipher from the `ctr` crate.
|
//! expanded key schedules are also zeroized by the RustCrypto backend.
|
||||||
//! The expanded key schedule lives inside that type and cannot be
|
//! - `AesCtr` uses the RustCrypto `zeroize` contract to clear its expanded
|
||||||
//! zeroized from outside. Callers that hold raw key material (e.g.
|
//! key schedule, counter, and buffered keystream on drop.
|
||||||
//! `HandshakeSuccess`, `ObfuscationParams`) are responsible for
|
//! - Callers that hold raw key material (e.g. `HandshakeSuccess`,
|
||||||
//! zeroizing their own copies.
|
//! `ObfuscationParams`) remain responsible for zeroizing their own copies.
|
||||||
|
|
||||||
#![allow(dead_code)]
|
#![allow(dead_code)]
|
||||||
|
|
||||||
@@ -19,24 +19,28 @@ use ctr::{
|
|||||||
Ctr128BE,
|
Ctr128BE,
|
||||||
cipher::{KeyIvInit, StreamCipher},
|
cipher::{KeyIvInit, StreamCipher},
|
||||||
};
|
};
|
||||||
use zeroize::Zeroize;
|
use zeroize::{Zeroize, ZeroizeOnDrop};
|
||||||
|
|
||||||
type Aes256Ctr = Ctr128BE<Aes256>;
|
type Aes256Ctr = Ctr128BE<Aes256>;
|
||||||
|
|
||||||
|
static_assertions::assert_impl_all!(Aes256: ZeroizeOnDrop);
|
||||||
|
static_assertions::assert_impl_all!(Aes256Ctr: ZeroizeOnDrop);
|
||||||
|
|
||||||
// ============= AES-256-CTR =============
|
// ============= AES-256-CTR =============
|
||||||
|
|
||||||
/// AES-256-CTR encryptor/decryptor
|
/// AES-256-CTR encryptor/decryptor
|
||||||
///
|
///
|
||||||
/// CTR mode is symmetric — encryption and decryption are the same operation.
|
/// CTR mode is symmetric — encryption and decryption are the same operation.
|
||||||
///
|
///
|
||||||
/// **Zeroize note:** The inner `Aes256Ctr` cipher state (expanded key schedule
|
/// **Zeroize note:** The inner `Aes256Ctr` zeroizes its expanded key schedule,
|
||||||
/// + counter) is opaque and cannot be zeroized. If you need to protect key
|
/// counter, and buffered keystream on drop. Callers remain responsible for
|
||||||
/// material, zeroize the `[u8; 32]` key and `u128` IV at the call site
|
/// zeroizing their own raw key and IV copies.
|
||||||
/// before dropping them.
|
|
||||||
pub struct AesCtr {
|
pub struct AesCtr {
|
||||||
cipher: Aes256Ctr,
|
cipher: Aes256Ctr,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl ZeroizeOnDrop for AesCtr {}
|
||||||
|
|
||||||
impl AesCtr {
|
impl AesCtr {
|
||||||
/// Create new AES-CTR cipher with key and IV
|
/// Create new AES-CTR cipher with key and IV
|
||||||
pub fn new(key: &[u8; 32], iv: u128) -> Self {
|
pub fn new(key: &[u8; 32], iv: u128) -> Self {
|
||||||
@@ -92,7 +96,7 @@ impl AesCtr {
|
|||||||
/// are different operations. This implementation handles CBC chaining
|
/// are different operations. This implementation handles CBC chaining
|
||||||
/// correctly across multiple blocks.
|
/// correctly across multiple blocks.
|
||||||
///
|
///
|
||||||
/// Key and IV are zeroized on drop.
|
/// Key, IV, and temporary expanded key schedules are zeroized on drop.
|
||||||
pub struct AesCbc {
|
pub struct AesCbc {
|
||||||
key: [u8; 32],
|
key: [u8; 32],
|
||||||
iv: [u8; 16],
|
iv: [u8; 16],
|
||||||
@@ -105,6 +109,8 @@ impl Drop for AesCbc {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl ZeroizeOnDrop for AesCbc {}
|
||||||
|
|
||||||
impl AesCbc {
|
impl AesCbc {
|
||||||
/// AES block size
|
/// AES block size
|
||||||
const BLOCK_SIZE: usize = 16;
|
const BLOCK_SIZE: usize = 16;
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
+20
-1
@@ -33,6 +33,9 @@ use crate::conntrack_control;
|
|||||||
use crate::crypto::SecureRandom;
|
use crate::crypto::SecureRandom;
|
||||||
use crate::ip_tracker::UserIpTracker;
|
use crate::ip_tracker::UserIpTracker;
|
||||||
use crate::network::probe::{decide_network_capabilities, log_probe_result, run_probe};
|
use crate::network::probe::{decide_network_capabilities, log_probe_result, run_probe};
|
||||||
|
use crate::proxy::direct_buffer_budget::{
|
||||||
|
DirectBufferBudget, resolve_direct_buffer_hard_limit, spawn_direct_buffer_budget_controller,
|
||||||
|
};
|
||||||
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
|
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
|
||||||
use crate::proxy::shared_state::ProxySharedState;
|
use crate::proxy::shared_state::ProxySharedState;
|
||||||
use crate::startup::{
|
use crate::startup::{
|
||||||
@@ -473,7 +476,16 @@ async fn run_telemt_core(
|
|||||||
config.network.dns_overrides.len()
|
config.network.dns_overrides.len()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
let shared_state = ProxySharedState::new();
|
let direct_buffer_hard_limit =
|
||||||
|
resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await;
|
||||||
|
let direct_buffer_budget = DirectBufferBudget::new(direct_buffer_hard_limit);
|
||||||
|
info!(
|
||||||
|
hard_limit_bytes = direct_buffer_hard_limit,
|
||||||
|
configured_override_bytes = config.general.direct_relay_buffer_budget_max_bytes,
|
||||||
|
"Direct relay buffer budget initialized"
|
||||||
|
);
|
||||||
|
let shared_state =
|
||||||
|
ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone());
|
||||||
shared_state.apply_user_enabled_config(&config.access.user_enabled);
|
shared_state.apply_user_enabled_config(&config.access.user_enabled);
|
||||||
shared_state.traffic_limiter.apply_policy(
|
shared_state.traffic_limiter.apply_policy(
|
||||||
config.access.user_rate_limits.clone(),
|
config.access.user_rate_limits.clone(),
|
||||||
@@ -882,6 +894,13 @@ async fn run_telemt_core(
|
|||||||
stats.clone(),
|
stats.clone(),
|
||||||
shared_state.clone(),
|
shared_state.clone(),
|
||||||
);
|
);
|
||||||
|
spawn_direct_buffer_budget_controller(
|
||||||
|
direct_buffer_budget,
|
||||||
|
buffer_pool.clone(),
|
||||||
|
stats.clone(),
|
||||||
|
shared_state.clone(),
|
||||||
|
config.server.max_connections,
|
||||||
|
);
|
||||||
|
|
||||||
let bound = listeners::bind_listeners(
|
let bound = listeners::bind_listeners(
|
||||||
&config,
|
&config,
|
||||||
|
|||||||
+151
@@ -595,6 +595,81 @@ async fn render_metrics(
|
|||||||
"telemt_buffer_pool_buffers_total{{kind=\"in_use\"}} {}",
|
"telemt_buffer_pool_buffers_total{{kind=\"in_use\"}} {}",
|
||||||
stats.get_buffer_pool_in_use_gauge()
|
stats.get_buffer_pool_in_use_gauge()
|
||||||
);
|
);
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"# HELP telemt_buffer_pool_events_total Buffer-pool allocation lifecycle events"
|
||||||
|
);
|
||||||
|
let _ = writeln!(out, "# TYPE telemt_buffer_pool_events_total counter");
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"telemt_buffer_pool_events_total{{event=\"replaced_nonstandard\"}} {}",
|
||||||
|
stats.get_buffer_pool_replaced_nonstandard_total()
|
||||||
|
);
|
||||||
|
|
||||||
|
let direct_budget = shared_state.direct_buffer_budget.snapshot();
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"# HELP telemt_direct_relay_buffer_budget_bytes Direct relay copy-buffer budget and memory inputs"
|
||||||
|
);
|
||||||
|
let _ = writeln!(out, "# TYPE telemt_direct_relay_buffer_budget_bytes gauge");
|
||||||
|
for (kind, value) in [
|
||||||
|
("hard_limit", direct_budget.hard_limit_bytes),
|
||||||
|
("target", direct_budget.target_bytes),
|
||||||
|
("reserved", direct_budget.reserved_bytes),
|
||||||
|
("memory_total", direct_budget.memory_total_bytes),
|
||||||
|
("memory_available", direct_budget.memory_available_bytes),
|
||||||
|
("process_rss", direct_budget.process_rss_bytes),
|
||||||
|
] {
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"telemt_direct_relay_buffer_budget_bytes{{kind=\"{}\"}} {}",
|
||||||
|
kind, value
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"# HELP telemt_direct_relay_buffer_budget_events_total Direct relay buffer-budget lifecycle events"
|
||||||
|
);
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"# TYPE telemt_direct_relay_buffer_budget_events_total counter"
|
||||||
|
);
|
||||||
|
for (result, value) in [
|
||||||
|
("promotion", direct_budget.promotion_total),
|
||||||
|
("promotion_denied", direct_budget.promotion_denied_total),
|
||||||
|
("minimum_fallback", direct_budget.minimum_fallback_total),
|
||||||
|
("admission_rejected", direct_budget.admission_rejected_total),
|
||||||
|
("quiet_demotion", direct_budget.quiet_demotion_total),
|
||||||
|
(
|
||||||
|
"write_pressure_demotion",
|
||||||
|
direct_budget.write_pressure_demotion_total,
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"global_pressure_demotion",
|
||||||
|
direct_budget.global_pressure_demotion_total,
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"telemt_direct_relay_buffer_budget_events_total{{result=\"{}\"}} {}",
|
||||||
|
result, value
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"# HELP telemt_direct_relay_buffer_sessions Current Direct relay sessions by adaptive tier"
|
||||||
|
);
|
||||||
|
let _ = writeln!(out, "# TYPE telemt_direct_relay_buffer_sessions gauge");
|
||||||
|
for (tier, value) in ["base", "tier1", "tier2", "tier3"]
|
||||||
|
.into_iter()
|
||||||
|
.zip(direct_budget.tier_sessions)
|
||||||
|
{
|
||||||
|
let _ = writeln!(
|
||||||
|
out,
|
||||||
|
"telemt_direct_relay_buffer_sessions{{tier=\"{}\"}} {}",
|
||||||
|
tier, value
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
let _ = writeln!(
|
let _ = writeln!(
|
||||||
out,
|
out,
|
||||||
@@ -2435,6 +2510,82 @@ 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"
|
||||||
|
|||||||
+137
-41
@@ -1,7 +1,4 @@
|
|||||||
#![allow(dead_code)]
|
// Adaptive buffer policy shared by active Direct relay sessions.
|
||||||
|
|
||||||
// Adaptive buffer policy is staged and retained for deterministic rollout.
|
|
||||||
// Keep definitions compiled for compatibility and security test scaffolding.
|
|
||||||
|
|
||||||
use dashmap::DashMap;
|
use dashmap::DashMap;
|
||||||
use std::cmp::max;
|
use std::cmp::max;
|
||||||
@@ -13,29 +10,42 @@ const PROFILE_TTL: Duration = Duration::from_secs(300);
|
|||||||
const THROUGHPUT_UP_BPS: f64 = 8_000_000.0;
|
const THROUGHPUT_UP_BPS: f64 = 8_000_000.0;
|
||||||
const THROUGHPUT_DOWN_BPS: f64 = 2_000_000.0;
|
const THROUGHPUT_DOWN_BPS: f64 = 2_000_000.0;
|
||||||
const RATIO_CONFIRM_THRESHOLD: f64 = 1.12;
|
const RATIO_CONFIRM_THRESHOLD: f64 = 1.12;
|
||||||
const TIER1_HOLD_TICKS: u32 = 8;
|
const TIER1_HOLD: Duration = Duration::from_secs(2);
|
||||||
const TIER2_HOLD_TICKS: u32 = 4;
|
const TIER2_HOLD: Duration = Duration::from_secs(1);
|
||||||
const QUIET_DEMOTE_TICKS: u32 = 480;
|
const QUIET_DEMOTE: Duration = Duration::from_secs(120);
|
||||||
const HARD_COOLDOWN_TICKS: u32 = 20;
|
const HARD_COOLDOWN: Duration = Duration::from_secs(5);
|
||||||
|
const SUSTAINED_PRESSURE_DEMOTE: Duration = Duration::from_secs(30);
|
||||||
|
const PRESSURE_DEMOTE_COOLDOWN: Duration = Duration::from_secs(60);
|
||||||
const HARD_PENDING_THRESHOLD: u32 = 3;
|
const HARD_PENDING_THRESHOLD: u32 = 3;
|
||||||
const HARD_PARTIAL_RATIO_THRESHOLD: f64 = 0.25;
|
const HARD_PARTIAL_RATIO_THRESHOLD: f64 = 0.25;
|
||||||
|
#[cfg(test)]
|
||||||
const DIRECT_C2S_CAP_BYTES: usize = 128 * 1024;
|
const DIRECT_C2S_CAP_BYTES: usize = 128 * 1024;
|
||||||
|
#[cfg(test)]
|
||||||
const DIRECT_S2C_CAP_BYTES: usize = 512 * 1024;
|
const DIRECT_S2C_CAP_BYTES: usize = 512 * 1024;
|
||||||
|
#[cfg(test)]
|
||||||
const ME_FRAMES_CAP: usize = 96;
|
const ME_FRAMES_CAP: usize = 96;
|
||||||
|
#[cfg(test)]
|
||||||
const ME_BYTES_CAP: usize = 384 * 1024;
|
const ME_BYTES_CAP: usize = 384 * 1024;
|
||||||
|
#[cfg(test)]
|
||||||
const ME_DELAY_MIN_US: u64 = 150;
|
const ME_DELAY_MIN_US: u64 = 150;
|
||||||
const MAX_USER_PROFILES_ENTRIES: usize = 50_000;
|
const MAX_USER_PROFILES_ENTRIES: usize = 50_000;
|
||||||
const MAX_USER_KEY_BYTES: usize = 512;
|
const MAX_USER_KEY_BYTES: usize = 512;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||||
|
/// Per-session Direct copy-buffer capacity tier.
|
||||||
pub enum AdaptiveTier {
|
pub enum AdaptiveTier {
|
||||||
|
/// Conservative baseline capacity.
|
||||||
Base = 0,
|
Base = 0,
|
||||||
|
/// First throughput promotion.
|
||||||
Tier1 = 1,
|
Tier1 = 1,
|
||||||
|
/// Sustained bidirectional pressure promotion.
|
||||||
Tier2 = 2,
|
Tier2 = 2,
|
||||||
|
/// Configured per-direction ceilings.
|
||||||
Tier3 = 3,
|
Tier3 = 3,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl AdaptiveTier {
|
impl AdaptiveTier {
|
||||||
|
/// Returns the next larger tier, saturating at `Tier3`.
|
||||||
pub fn promote(self) -> Self {
|
pub fn promote(self) -> Self {
|
||||||
match self {
|
match self {
|
||||||
Self::Base => Self::Tier1,
|
Self::Base => Self::Tier1,
|
||||||
@@ -45,6 +55,7 @@ impl AdaptiveTier {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the next smaller tier, saturating at `Base`.
|
||||||
pub fn demote(self) -> Self {
|
pub fn demote(self) -> Self {
|
||||||
match self {
|
match self {
|
||||||
Self::Base => Self::Base,
|
Self::Base => Self::Base,
|
||||||
@@ -54,6 +65,7 @@ impl AdaptiveTier {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
fn ratio(self) -> (usize, usize) {
|
fn ratio(self) -> (usize, usize) {
|
||||||
match self {
|
match self {
|
||||||
Self::Base => (1, 1),
|
Self::Base => (1, 1),
|
||||||
@@ -63,49 +75,70 @@ impl AdaptiveTier {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the stable numeric tier used by bounded metrics.
|
||||||
pub fn as_u8(self) -> u8 {
|
pub fn as_u8(self) -> u8 {
|
||||||
self as u8
|
self as u8
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
/// Signal that caused an accepted adaptive tier transition.
|
||||||
pub enum TierTransitionReason {
|
pub enum TierTransitionReason {
|
||||||
|
/// Sustained throughput and directional ratio confirmation.
|
||||||
SoftConfirmed,
|
SoftConfirmed,
|
||||||
|
/// Short pending or partial-write pressure burst.
|
||||||
HardPressure,
|
HardPressure,
|
||||||
|
/// Sustained low-throughput period.
|
||||||
QuietDemotion,
|
QuietDemotion,
|
||||||
|
/// Sustained pending or partial-write pressure.
|
||||||
|
SustainedWritePressure,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
/// Proposed transition emitted by the per-session controller.
|
||||||
pub struct TierTransition {
|
pub struct TierTransition {
|
||||||
|
/// Tier active before the observation.
|
||||||
pub from: AdaptiveTier,
|
pub from: AdaptiveTier,
|
||||||
|
/// Tier requested after the observation.
|
||||||
pub to: AdaptiveTier,
|
pub to: AdaptiveTier,
|
||||||
|
/// Pressure or throughput condition that requested the transition.
|
||||||
pub reason: TierTransitionReason,
|
pub reason: TierTransitionReason,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, Default)]
|
#[derive(Debug, Clone, Copy, Default)]
|
||||||
|
/// Directional byte and write-pressure deltas for one observation period.
|
||||||
pub struct RelaySignalSample {
|
pub struct RelaySignalSample {
|
||||||
|
/// Client-to-DC bytes copied during the period.
|
||||||
pub c2s_bytes: u64,
|
pub c2s_bytes: u64,
|
||||||
|
/// Bytes offered to DC-to-client writes during the period.
|
||||||
pub s2c_requested_bytes: u64,
|
pub s2c_requested_bytes: u64,
|
||||||
|
/// Bytes accepted by DC-to-client writes during the period.
|
||||||
pub s2c_written_bytes: u64,
|
pub s2c_written_bytes: u64,
|
||||||
|
/// Successful DC-to-client write operations during the period.
|
||||||
pub s2c_write_ops: u64,
|
pub s2c_write_ops: u64,
|
||||||
|
/// Partial DC-to-client write operations during the period.
|
||||||
pub s2c_partial_writes: u64,
|
pub s2c_partial_writes: u64,
|
||||||
|
/// Consecutive pending DC-to-client writes at sample time.
|
||||||
pub s2c_consecutive_pending_writes: u32,
|
pub s2c_consecutive_pending_writes: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy)]
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
/// Stateful hysteresis controller for one active Direct session.
|
||||||
pub struct SessionAdaptiveController {
|
pub struct SessionAdaptiveController {
|
||||||
tier: AdaptiveTier,
|
tier: AdaptiveTier,
|
||||||
max_tier_seen: AdaptiveTier,
|
max_tier_seen: AdaptiveTier,
|
||||||
throughput_ema_bps: f64,
|
throughput_ema_bps: f64,
|
||||||
incoming_ema_bps: f64,
|
incoming_ema_bps: f64,
|
||||||
outgoing_ema_bps: f64,
|
outgoing_ema_bps: f64,
|
||||||
tier1_hold_ticks: u32,
|
tier1_hold: Duration,
|
||||||
tier2_hold_ticks: u32,
|
tier2_hold: Duration,
|
||||||
quiet_ticks: u32,
|
quiet: Duration,
|
||||||
hard_cooldown_ticks: u32,
|
hard_cooldown: Duration,
|
||||||
|
sustained_pressure: Duration,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl SessionAdaptiveController {
|
impl SessionAdaptiveController {
|
||||||
|
/// Creates a controller at the tier whose memory reservation was accepted.
|
||||||
pub fn new(initial_tier: AdaptiveTier) -> Self {
|
pub fn new(initial_tier: AdaptiveTier) -> Self {
|
||||||
Self {
|
Self {
|
||||||
tier: initial_tier,
|
tier: initial_tier,
|
||||||
@@ -113,25 +146,33 @@ impl SessionAdaptiveController {
|
|||||||
throughput_ema_bps: 0.0,
|
throughput_ema_bps: 0.0,
|
||||||
incoming_ema_bps: 0.0,
|
incoming_ema_bps: 0.0,
|
||||||
outgoing_ema_bps: 0.0,
|
outgoing_ema_bps: 0.0,
|
||||||
tier1_hold_ticks: 0,
|
tier1_hold: Duration::ZERO,
|
||||||
tier2_hold_ticks: 0,
|
tier2_hold: Duration::ZERO,
|
||||||
quiet_ticks: 0,
|
quiet: Duration::ZERO,
|
||||||
hard_cooldown_ticks: 0,
|
hard_cooldown: Duration::ZERO,
|
||||||
|
sustained_pressure: Duration::ZERO,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the highest tier proposed by controller observations.
|
||||||
|
#[allow(dead_code)]
|
||||||
pub fn max_tier_seen(&self) -> AdaptiveTier {
|
pub fn max_tier_seen(&self) -> AdaptiveTier {
|
||||||
self.max_tier_seen
|
self.max_tier_seen
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the controller's current logical tier.
|
||||||
|
pub fn tier(&self) -> AdaptiveTier {
|
||||||
|
self.tier
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Observes one period and returns at most one hysteresis-controlled transition.
|
||||||
pub fn observe(&mut self, sample: RelaySignalSample, tick_secs: f64) -> Option<TierTransition> {
|
pub fn observe(&mut self, sample: RelaySignalSample, tick_secs: f64) -> Option<TierTransition> {
|
||||||
if tick_secs <= f64::EPSILON {
|
if tick_secs <= f64::EPSILON {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
if self.hard_cooldown_ticks > 0 {
|
let tick = Duration::from_secs_f64(tick_secs);
|
||||||
self.hard_cooldown_ticks -= 1;
|
self.hard_cooldown = self.hard_cooldown.saturating_sub(tick);
|
||||||
}
|
|
||||||
|
|
||||||
let c2s_bps = (sample.c2s_bytes as f64 * 8.0) / tick_secs;
|
let c2s_bps = (sample.c2s_bytes as f64 * 8.0) / tick_secs;
|
||||||
let incoming_bps = (sample.s2c_requested_bytes as f64 * 8.0) / tick_secs;
|
let incoming_bps = (sample.s2c_requested_bytes as f64 * 8.0) / tick_secs;
|
||||||
@@ -144,9 +185,9 @@ impl SessionAdaptiveController {
|
|||||||
|
|
||||||
let tier1_now = self.throughput_ema_bps >= THROUGHPUT_UP_BPS;
|
let tier1_now = self.throughput_ema_bps >= THROUGHPUT_UP_BPS;
|
||||||
if tier1_now {
|
if tier1_now {
|
||||||
self.tier1_hold_ticks = self.tier1_hold_ticks.saturating_add(1);
|
self.tier1_hold = self.tier1_hold.saturating_add(tick);
|
||||||
} else {
|
} else {
|
||||||
self.tier1_hold_ticks = 0;
|
self.tier1_hold = Duration::ZERO;
|
||||||
}
|
}
|
||||||
|
|
||||||
let ratio = if self.outgoing_ema_bps <= f64::EPSILON {
|
let ratio = if self.outgoing_ema_bps <= f64::EPSILON {
|
||||||
@@ -156,9 +197,9 @@ impl SessionAdaptiveController {
|
|||||||
};
|
};
|
||||||
let tier2_now = ratio >= RATIO_CONFIRM_THRESHOLD;
|
let tier2_now = ratio >= RATIO_CONFIRM_THRESHOLD;
|
||||||
if tier2_now {
|
if tier2_now {
|
||||||
self.tier2_hold_ticks = self.tier2_hold_ticks.saturating_add(1);
|
self.tier2_hold = self.tier2_hold.saturating_add(tick);
|
||||||
} else {
|
} else {
|
||||||
self.tier2_hold_ticks = 0;
|
self.tier2_hold = Duration::ZERO;
|
||||||
}
|
}
|
||||||
|
|
||||||
let partial_ratio = if sample.s2c_write_ops == 0 {
|
let partial_ratio = if sample.s2c_write_ops == 0 {
|
||||||
@@ -169,24 +210,37 @@ impl SessionAdaptiveController {
|
|||||||
let hard_now = sample.s2c_consecutive_pending_writes >= HARD_PENDING_THRESHOLD
|
let hard_now = sample.s2c_consecutive_pending_writes >= HARD_PENDING_THRESHOLD
|
||||||
|| partial_ratio >= HARD_PARTIAL_RATIO_THRESHOLD;
|
|| partial_ratio >= HARD_PARTIAL_RATIO_THRESHOLD;
|
||||||
|
|
||||||
if hard_now && self.hard_cooldown_ticks == 0 {
|
if hard_now {
|
||||||
return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN_TICKS);
|
self.sustained_pressure = self.sustained_pressure.saturating_add(tick);
|
||||||
|
if self.sustained_pressure >= SUSTAINED_PRESSURE_DEMOTE {
|
||||||
|
self.sustained_pressure = Duration::ZERO;
|
||||||
|
return self.demote(
|
||||||
|
TierTransitionReason::SustainedWritePressure,
|
||||||
|
PRESSURE_DEMOTE_COOLDOWN,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
self.sustained_pressure = Duration::ZERO;
|
||||||
}
|
}
|
||||||
|
|
||||||
if self.tier1_hold_ticks >= TIER1_HOLD_TICKS && self.tier2_hold_ticks >= TIER2_HOLD_TICKS {
|
if hard_now && self.hard_cooldown.is_zero() {
|
||||||
return self.promote(TierTransitionReason::SoftConfirmed, 0);
|
return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN);
|
||||||
|
}
|
||||||
|
|
||||||
|
if self.tier1_hold >= TIER1_HOLD && self.tier2_hold >= TIER2_HOLD {
|
||||||
|
return self.promote(TierTransitionReason::SoftConfirmed, Duration::ZERO);
|
||||||
}
|
}
|
||||||
|
|
||||||
let demote_candidate =
|
let demote_candidate =
|
||||||
self.throughput_ema_bps < THROUGHPUT_DOWN_BPS && !tier2_now && !hard_now;
|
self.throughput_ema_bps < THROUGHPUT_DOWN_BPS && !tier2_now && !hard_now;
|
||||||
if demote_candidate {
|
if demote_candidate {
|
||||||
self.quiet_ticks = self.quiet_ticks.saturating_add(1);
|
self.quiet = self.quiet.saturating_add(tick);
|
||||||
if self.quiet_ticks >= QUIET_DEMOTE_TICKS {
|
if self.quiet >= QUIET_DEMOTE {
|
||||||
self.quiet_ticks = 0;
|
self.quiet = Duration::ZERO;
|
||||||
return self.demote(TierTransitionReason::QuietDemotion);
|
return self.demote(TierTransitionReason::QuietDemotion, Duration::ZERO);
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
self.quiet_ticks = 0;
|
self.quiet = Duration::ZERO;
|
||||||
}
|
}
|
||||||
|
|
||||||
None
|
None
|
||||||
@@ -195,7 +249,7 @@ impl SessionAdaptiveController {
|
|||||||
fn promote(
|
fn promote(
|
||||||
&mut self,
|
&mut self,
|
||||||
reason: TierTransitionReason,
|
reason: TierTransitionReason,
|
||||||
hard_cooldown_ticks: u32,
|
hard_cooldown: Duration,
|
||||||
) -> Option<TierTransition> {
|
) -> Option<TierTransition> {
|
||||||
let from = self.tier;
|
let from = self.tier;
|
||||||
let to = from.promote();
|
let to = from.promote();
|
||||||
@@ -204,22 +258,27 @@ impl SessionAdaptiveController {
|
|||||||
}
|
}
|
||||||
self.tier = to;
|
self.tier = to;
|
||||||
self.max_tier_seen = max(self.max_tier_seen, to);
|
self.max_tier_seen = max(self.max_tier_seen, to);
|
||||||
self.hard_cooldown_ticks = hard_cooldown_ticks;
|
self.hard_cooldown = hard_cooldown;
|
||||||
self.tier1_hold_ticks = 0;
|
self.tier1_hold = Duration::ZERO;
|
||||||
self.tier2_hold_ticks = 0;
|
self.tier2_hold = Duration::ZERO;
|
||||||
self.quiet_ticks = 0;
|
self.quiet = Duration::ZERO;
|
||||||
Some(TierTransition { from, to, reason })
|
Some(TierTransition { from, to, reason })
|
||||||
}
|
}
|
||||||
|
|
||||||
fn demote(&mut self, reason: TierTransitionReason) -> Option<TierTransition> {
|
fn demote(
|
||||||
|
&mut self,
|
||||||
|
reason: TierTransitionReason,
|
||||||
|
hard_cooldown: Duration,
|
||||||
|
) -> Option<TierTransition> {
|
||||||
let from = self.tier;
|
let from = self.tier;
|
||||||
let to = from.demote();
|
let to = from.demote();
|
||||||
if from == to {
|
if from == to {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
self.tier = to;
|
self.tier = to;
|
||||||
self.tier1_hold_ticks = 0;
|
self.hard_cooldown = hard_cooldown;
|
||||||
self.tier2_hold_ticks = 0;
|
self.tier1_hold = Duration::ZERO;
|
||||||
|
self.tier2_hold = Duration::ZERO;
|
||||||
Some(TierTransition { from, to, reason })
|
Some(TierTransition { from, to, reason })
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -235,6 +294,8 @@ fn profiles() -> &'static DashMap<String, UserAdaptiveProfile> {
|
|||||||
USER_PROFILES.get_or_init(DashMap::new)
|
USER_PROFILES.get_or_init(DashMap::new)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns a fresh user's recent successful Direct tier, or `Base` when stale.
|
||||||
|
#[allow(dead_code)]
|
||||||
pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
|
pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
|
||||||
if user.len() > MAX_USER_KEY_BYTES {
|
if user.len() > MAX_USER_KEY_BYTES {
|
||||||
return AdaptiveTier::Base;
|
return AdaptiveTier::Base;
|
||||||
@@ -253,6 +314,8 @@ pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
|
|||||||
AdaptiveTier::Base
|
AdaptiveTier::Base
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Records the highest successfully allocated tier for bounded session seeding.
|
||||||
|
#[allow(dead_code)]
|
||||||
pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
|
pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
|
||||||
if user.len() > MAX_USER_KEY_BYTES {
|
if user.len() > MAX_USER_KEY_BYTES {
|
||||||
return;
|
return;
|
||||||
@@ -282,6 +345,8 @@ pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
/// Returns the legacy staged scaling policy retained by security fixtures.
|
||||||
pub fn direct_copy_buffers_for_tier(
|
pub fn direct_copy_buffers_for_tier(
|
||||||
tier: AdaptiveTier,
|
tier: AdaptiveTier,
|
||||||
base_c2s: usize,
|
base_c2s: usize,
|
||||||
@@ -294,6 +359,32 @@ pub fn direct_copy_buffers_for_tier(
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Maps an adaptive tier to independent capacities within configured ceilings.
|
||||||
|
pub(crate) fn direct_copy_buffers_for_tier_with_ceilings(
|
||||||
|
tier: AdaptiveTier,
|
||||||
|
base_c2s: usize,
|
||||||
|
base_s2c: usize,
|
||||||
|
ceiling_c2s: usize,
|
||||||
|
ceiling_s2c: usize,
|
||||||
|
) -> (usize, usize) {
|
||||||
|
(
|
||||||
|
direct_direction_size(tier, base_c2s, ceiling_c2s),
|
||||||
|
direct_direction_size(tier, base_s2c, ceiling_s2c),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn direct_direction_size(tier: AdaptiveTier, base: usize, ceiling: usize) -> usize {
|
||||||
|
let target = match tier {
|
||||||
|
AdaptiveTier::Base => base,
|
||||||
|
AdaptiveTier::Tier1 => ceiling / 4,
|
||||||
|
AdaptiveTier::Tier2 => ceiling / 2,
|
||||||
|
AdaptiveTier::Tier3 => ceiling,
|
||||||
|
};
|
||||||
|
target.max(base).min(ceiling.max(base)).max(1)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
/// Returns the staged Middle-End flush policy retained by security fixtures.
|
||||||
pub fn me_flush_policy_for_tier(
|
pub fn me_flush_policy_for_tier(
|
||||||
tier: AdaptiveTier,
|
tier: AdaptiveTier,
|
||||||
base_frames: usize,
|
base_frames: usize,
|
||||||
@@ -323,6 +414,7 @@ fn ema(prev: f64, value: f64) -> f64 {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
fn scale(base: usize, numerator: usize, denominator: usize, cap: usize) -> usize {
|
fn scale(base: usize, numerator: usize, denominator: usize, cap: usize) -> usize {
|
||||||
let scaled = base
|
let scaled = base
|
||||||
.saturating_mul(numerator)
|
.saturating_mul(numerator)
|
||||||
@@ -338,6 +430,10 @@ mod adaptive_buffers_security_tests;
|
|||||||
#[path = "tests/adaptive_buffers_record_race_security_tests.rs"]
|
#[path = "tests/adaptive_buffers_record_race_security_tests.rs"]
|
||||||
mod adaptive_buffers_record_race_security_tests;
|
mod adaptive_buffers_record_race_security_tests;
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "tests/adaptive_direct_budget_policy_tests.rs"]
|
||||||
|
mod adaptive_direct_budget_policy_tests;
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -396,7 +492,7 @@ mod tests {
|
|||||||
fn test_quiet_demotion_is_slow_and_stepwise() {
|
fn test_quiet_demotion_is_slow_and_stepwise() {
|
||||||
let mut ctrl = SessionAdaptiveController::new(AdaptiveTier::Tier2);
|
let mut ctrl = SessionAdaptiveController::new(AdaptiveTier::Tier2);
|
||||||
let mut demotion = None;
|
let mut demotion = None;
|
||||||
for _ in 0..QUIET_DEMOTE_TICKS {
|
for _ in 0..480 {
|
||||||
demotion = ctrl.observe(sample(1, 1, 1, 1, 0, 0), 0.25);
|
demotion = ctrl.observe(sample(1, 1, 1, 1, 0, 0), 0.25);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,570 @@
|
|||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use tokio::sync::watch;
|
||||||
|
|
||||||
|
use crate::stats::Stats;
|
||||||
|
use crate::stream::BufferPool;
|
||||||
|
|
||||||
|
use super::shared_state::ProxySharedState;
|
||||||
|
|
||||||
|
/// Accounting granularity for process-wide Direct copy-buffer reservations.
|
||||||
|
pub(crate) const DIRECT_BUFFER_UNIT_BYTES: usize = 4 * 1024;
|
||||||
|
/// Minimum client-to-DC copy-buffer capacity for one Direct session.
|
||||||
|
pub(crate) const DIRECT_BASE_C2S_BYTES: usize = 4 * 1024;
|
||||||
|
/// Minimum DC-to-client copy-buffer capacity for one Direct session.
|
||||||
|
pub(crate) const DIRECT_BASE_S2C_BYTES: usize = 8 * 1024;
|
||||||
|
|
||||||
|
const AUTO_HARD_MIN_BYTES: usize = 64 * 1024 * 1024;
|
||||||
|
const AUTO_HARD_MAX_BYTES: usize = 2 * 1024 * 1024 * 1024;
|
||||||
|
const AUTO_HARD_FALLBACK_BYTES: usize = 512 * 1024 * 1024;
|
||||||
|
const TARGET_FLOOR_MIN_BYTES: usize = 16 * 1024 * 1024;
|
||||||
|
const CONTROL_INTERVAL: Duration = Duration::from_secs(1);
|
||||||
|
const HEALTHY_RECOVERY_SAMPLES: u8 = 30;
|
||||||
|
const BUFFER_POOL_TRIM_LOW_WATERMARK: usize = 64;
|
||||||
|
const BUFFER_POOL_TRIM_HIGH_WATERMARK: usize = 128;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, Default)]
|
||||||
|
/// Lock-free observability snapshot of the Direct copy-buffer envelope.
|
||||||
|
pub(crate) struct DirectBufferBudgetSnapshot {
|
||||||
|
/// Absolute process-wide copy-buffer ceiling.
|
||||||
|
pub(crate) hard_limit_bytes: u64,
|
||||||
|
/// Current pressure-adjusted promotion target.
|
||||||
|
pub(crate) target_bytes: u64,
|
||||||
|
/// Bytes currently covered by active session leases.
|
||||||
|
pub(crate) reserved_bytes: u64,
|
||||||
|
/// Effective host or cgroup memory limit.
|
||||||
|
pub(crate) memory_total_bytes: u64,
|
||||||
|
/// Effective host or cgroup memory headroom.
|
||||||
|
pub(crate) memory_available_bytes: u64,
|
||||||
|
/// Current process resident set size.
|
||||||
|
pub(crate) process_rss_bytes: u64,
|
||||||
|
/// Successful tier growth reservations.
|
||||||
|
pub(crate) promotion_total: u64,
|
||||||
|
/// Tier growth attempts rejected by the adaptive target.
|
||||||
|
pub(crate) promotion_denied_total: u64,
|
||||||
|
/// Sessions admitted at minimum size above the adaptive target.
|
||||||
|
pub(crate) minimum_fallback_total: u64,
|
||||||
|
/// Sessions rejected by the absolute ceiling.
|
||||||
|
pub(crate) admission_rejected_total: u64,
|
||||||
|
/// Quiet-period tier reductions.
|
||||||
|
pub(crate) quiet_demotion_total: u64,
|
||||||
|
/// Sustained write-pressure tier reductions.
|
||||||
|
pub(crate) write_pressure_demotion_total: u64,
|
||||||
|
/// Process-wide pressure tier reductions.
|
||||||
|
pub(crate) global_pressure_demotion_total: u64,
|
||||||
|
/// Current sessions for Base through Tier3.
|
||||||
|
pub(crate) tier_sessions: [u64; 4],
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, Default)]
|
||||||
|
struct SystemMemorySample {
|
||||||
|
total_bytes: u64,
|
||||||
|
available_bytes: u64,
|
||||||
|
process_rss_bytes: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Process-wide hard envelope and adaptive target for Direct copy buffers.
|
||||||
|
pub(crate) struct DirectBufferBudget {
|
||||||
|
hard_limit_bytes: u64,
|
||||||
|
target_bytes: AtomicU64,
|
||||||
|
reserved_bytes: AtomicU64,
|
||||||
|
pressure_generation: AtomicU64,
|
||||||
|
pressure_tx: watch::Sender<u64>,
|
||||||
|
memory_total_bytes: AtomicU64,
|
||||||
|
memory_available_bytes: AtomicU64,
|
||||||
|
process_rss_bytes: AtomicU64,
|
||||||
|
promotion_total: AtomicU64,
|
||||||
|
promotion_denied_total: AtomicU64,
|
||||||
|
minimum_fallback_total: AtomicU64,
|
||||||
|
admission_rejected_total: AtomicU64,
|
||||||
|
quiet_demotion_total: AtomicU64,
|
||||||
|
write_pressure_demotion_total: AtomicU64,
|
||||||
|
global_pressure_demotion_total: AtomicU64,
|
||||||
|
tier_sessions: [AtomicU64; 4],
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DirectBufferBudget {
|
||||||
|
/// Creates an envelope with a fixed absolute ceiling.
|
||||||
|
pub(crate) fn new(hard_limit_bytes: usize) -> Arc<Self> {
|
||||||
|
let hard_limit_bytes = align_down(hard_limit_bytes.max(DIRECT_BUFFER_UNIT_BYTES)) as u64;
|
||||||
|
let (pressure_tx, _) = watch::channel(0);
|
||||||
|
Arc::new(Self {
|
||||||
|
hard_limit_bytes,
|
||||||
|
target_bytes: AtomicU64::new(hard_limit_bytes),
|
||||||
|
reserved_bytes: AtomicU64::new(0),
|
||||||
|
pressure_generation: AtomicU64::new(0),
|
||||||
|
pressure_tx,
|
||||||
|
memory_total_bytes: AtomicU64::new(0),
|
||||||
|
memory_available_bytes: AtomicU64::new(0),
|
||||||
|
process_rss_bytes: AtomicU64::new(0),
|
||||||
|
promotion_total: AtomicU64::new(0),
|
||||||
|
promotion_denied_total: AtomicU64::new(0),
|
||||||
|
minimum_fallback_total: AtomicU64::new(0),
|
||||||
|
admission_rejected_total: AtomicU64::new(0),
|
||||||
|
quiet_demotion_total: AtomicU64::new(0),
|
||||||
|
write_pressure_demotion_total: AtomicU64::new(0),
|
||||||
|
global_pressure_demotion_total: AtomicU64::new(0),
|
||||||
|
tier_sessions: std::array::from_fn(|_| AtomicU64::new(0)),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns the current pressure-adjusted reservation target.
|
||||||
|
pub(crate) fn target_bytes(&self) -> usize {
|
||||||
|
self.target_bytes.load(Ordering::Relaxed) as usize
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Subscribes to target reductions that require prompt session demotion.
|
||||||
|
pub(crate) fn subscribe_pressure(&self) -> watch::Receiver<u64> {
|
||||||
|
self.pressure_tx.subscribe()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Reserves bytes against either the adaptive target or the absolute ceiling.
|
||||||
|
pub(crate) fn try_reserve(
|
||||||
|
self: &Arc<Self>,
|
||||||
|
bytes: usize,
|
||||||
|
allow_above_target: bool,
|
||||||
|
) -> Option<DirectBufferLease> {
|
||||||
|
let bytes = align_up(bytes) as u64;
|
||||||
|
let limit = if allow_above_target {
|
||||||
|
self.hard_limit_bytes
|
||||||
|
} else {
|
||||||
|
self.target_bytes
|
||||||
|
.load(Ordering::Relaxed)
|
||||||
|
.min(self.hard_limit_bytes)
|
||||||
|
};
|
||||||
|
if !self.try_add_reserved(bytes, limit) {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
self.tier_sessions[0].fetch_add(1, Ordering::Relaxed);
|
||||||
|
Some(DirectBufferLease {
|
||||||
|
budget: Arc::clone(self),
|
||||||
|
reserved_bytes: bytes,
|
||||||
|
tier: 0,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn try_add_reserved(&self, bytes: u64, limit: u64) -> bool {
|
||||||
|
let mut current = self.reserved_bytes.load(Ordering::Acquire);
|
||||||
|
loop {
|
||||||
|
if bytes > limit.saturating_sub(current) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
match self.reserved_bytes.compare_exchange_weak(
|
||||||
|
current,
|
||||||
|
current + bytes,
|
||||||
|
Ordering::AcqRel,
|
||||||
|
Ordering::Acquire,
|
||||||
|
) {
|
||||||
|
Ok(_) => return true,
|
||||||
|
Err(observed) => current = observed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn target_floor_bytes(&self) -> u64 {
|
||||||
|
(self.hard_limit_bytes / 8)
|
||||||
|
.max(TARGET_FLOOR_MIN_BYTES as u64)
|
||||||
|
.min(self.hard_limit_bytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_target_bytes(&self, target: u64) {
|
||||||
|
let target =
|
||||||
|
align_down(target.clamp(self.target_floor_bytes(), self.hard_limit_bytes) as usize)
|
||||||
|
as u64;
|
||||||
|
let previous = self.target_bytes.swap(target, Ordering::AcqRel);
|
||||||
|
if target < previous {
|
||||||
|
let generation = self
|
||||||
|
.pressure_generation
|
||||||
|
.fetch_add(1, Ordering::AcqRel)
|
||||||
|
.wrapping_add(1);
|
||||||
|
self.pressure_tx.send_replace(generation);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn update_system_sample(&self, sample: SystemMemorySample) {
|
||||||
|
self.memory_total_bytes
|
||||||
|
.store(sample.total_bytes, Ordering::Relaxed);
|
||||||
|
self.memory_available_bytes
|
||||||
|
.store(sample.available_bytes, Ordering::Relaxed);
|
||||||
|
self.process_rss_bytes
|
||||||
|
.store(sample.process_rss_bytes, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Records a session that had to bypass the adaptive target at minimum size.
|
||||||
|
pub(crate) fn increment_minimum_fallback(&self) {
|
||||||
|
self.minimum_fallback_total.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Records a session rejected because the absolute ceiling was exhausted.
|
||||||
|
pub(crate) fn increment_admission_rejected(&self) {
|
||||||
|
self.admission_rejected_total
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Records a tier reduction after sustained low throughput.
|
||||||
|
pub(crate) fn increment_quiet_demotion(&self) {
|
||||||
|
self.quiet_demotion_total.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Records a tier reduction after sustained partial or pending writes.
|
||||||
|
pub(crate) fn increment_write_pressure_demotion(&self) {
|
||||||
|
self.write_pressure_demotion_total
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Records a tier reduction requested by the process-wide controller.
|
||||||
|
pub(crate) fn increment_global_pressure_demotion(&self) {
|
||||||
|
self.global_pressure_demotion_total
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Captures all bounded metrics without allocating or locking.
|
||||||
|
pub(crate) fn snapshot(&self) -> DirectBufferBudgetSnapshot {
|
||||||
|
DirectBufferBudgetSnapshot {
|
||||||
|
hard_limit_bytes: self.hard_limit_bytes,
|
||||||
|
target_bytes: self.target_bytes.load(Ordering::Relaxed),
|
||||||
|
reserved_bytes: self.reserved_bytes.load(Ordering::Relaxed),
|
||||||
|
memory_total_bytes: self.memory_total_bytes.load(Ordering::Relaxed),
|
||||||
|
memory_available_bytes: self.memory_available_bytes.load(Ordering::Relaxed),
|
||||||
|
process_rss_bytes: self.process_rss_bytes.load(Ordering::Relaxed),
|
||||||
|
promotion_total: self.promotion_total.load(Ordering::Relaxed),
|
||||||
|
promotion_denied_total: self.promotion_denied_total.load(Ordering::Relaxed),
|
||||||
|
minimum_fallback_total: self.minimum_fallback_total.load(Ordering::Relaxed),
|
||||||
|
admission_rejected_total: self.admission_rejected_total.load(Ordering::Relaxed),
|
||||||
|
quiet_demotion_total: self.quiet_demotion_total.load(Ordering::Relaxed),
|
||||||
|
write_pressure_demotion_total: self
|
||||||
|
.write_pressure_demotion_total
|
||||||
|
.load(Ordering::Relaxed),
|
||||||
|
global_pressure_demotion_total: self
|
||||||
|
.global_pressure_demotion_total
|
||||||
|
.load(Ordering::Relaxed),
|
||||||
|
tier_sessions: std::array::from_fn(|index| {
|
||||||
|
self.tier_sessions[index].load(Ordering::Relaxed)
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns the conservative ceiling used when memory discovery is unavailable.
|
||||||
|
pub(crate) fn fallback_direct_buffer_hard_limit() -> usize {
|
||||||
|
AUTO_HARD_FALLBACK_BYTES
|
||||||
|
}
|
||||||
|
|
||||||
|
/// RAII ownership of all copy-buffer bytes retained by one Direct session.
|
||||||
|
pub(crate) struct DirectBufferLease {
|
||||||
|
budget: Arc<DirectBufferBudget>,
|
||||||
|
reserved_bytes: u64,
|
||||||
|
tier: usize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl DirectBufferLease {
|
||||||
|
/// Returns the currently covered allocation rounded to accounting units.
|
||||||
|
pub(crate) fn reserved_bytes(&self) -> usize {
|
||||||
|
self.reserved_bytes as usize
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Attempts to cover a larger tier before its buffers are resized.
|
||||||
|
pub(crate) fn try_grow_to(&mut self, bytes: usize) -> bool {
|
||||||
|
let bytes = align_up(bytes) as u64;
|
||||||
|
if bytes <= self.reserved_bytes {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
let delta = bytes - self.reserved_bytes;
|
||||||
|
let limit = self
|
||||||
|
.budget
|
||||||
|
.target_bytes
|
||||||
|
.load(Ordering::Relaxed)
|
||||||
|
.min(self.budget.hard_limit_bytes);
|
||||||
|
if !self.budget.try_add_reserved(delta, limit) {
|
||||||
|
self.budget
|
||||||
|
.promotion_denied_total
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
self.reserved_bytes = bytes;
|
||||||
|
self.budget.promotion_total.fetch_add(1, Ordering::Relaxed);
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Releases bytes only after both directional buffers report smaller coverage.
|
||||||
|
pub(crate) fn shrink_to(&mut self, bytes: usize) {
|
||||||
|
let bytes = align_up(bytes) as u64;
|
||||||
|
if bytes >= self.reserved_bytes {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let released = self.reserved_bytes - bytes;
|
||||||
|
self.reserved_bytes = bytes;
|
||||||
|
self.budget
|
||||||
|
.reserved_bytes
|
||||||
|
.fetch_sub(released, Ordering::AcqRel);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Updates bounded per-tier session gauges for an accepted transition.
|
||||||
|
pub(crate) fn set_tier(&mut self, tier: usize) {
|
||||||
|
let tier = tier.min(self.budget.tier_sessions.len() - 1);
|
||||||
|
if tier == self.tier {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
decrement_saturating(&self.budget.tier_sessions[self.tier]);
|
||||||
|
self.budget.tier_sessions[tier].fetch_add(1, Ordering::Relaxed);
|
||||||
|
self.tier = tier;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for DirectBufferLease {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
self.budget
|
||||||
|
.reserved_bytes
|
||||||
|
.fetch_sub(self.reserved_bytes, Ordering::AcqRel);
|
||||||
|
decrement_saturating(&self.budget.tier_sessions[self.tier]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolves the startup hard ceiling from config, cgroup, and host memory.
|
||||||
|
pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize {
|
||||||
|
if configured != 0 {
|
||||||
|
return align_down(configured);
|
||||||
|
}
|
||||||
|
let sample = read_system_memory_sample().await;
|
||||||
|
if sample.total_bytes == 0 {
|
||||||
|
return AUTO_HARD_FALLBACK_BYTES;
|
||||||
|
}
|
||||||
|
let derived = (sample.total_bytes / 4)
|
||||||
|
.clamp(AUTO_HARD_MIN_BYTES as u64, AUTO_HARD_MAX_BYTES as u64)
|
||||||
|
.min(sample.total_bytes);
|
||||||
|
align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Starts the single control-plane task for Direct budget and shared pool pressure.
|
||||||
|
pub(crate) fn spawn_direct_buffer_budget_controller(
|
||||||
|
budget: Arc<DirectBufferBudget>,
|
||||||
|
buffer_pool: Arc<BufferPool>,
|
||||||
|
stats: Arc<Stats>,
|
||||||
|
shared: Arc<ProxySharedState>,
|
||||||
|
max_connections: u32,
|
||||||
|
) {
|
||||||
|
tokio::spawn(async move {
|
||||||
|
let mut interval = tokio::time::interval(CONTROL_INTERVAL);
|
||||||
|
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||||
|
let mut healthy_streak = 0u8;
|
||||||
|
let mut previous_denied = 0u64;
|
||||||
|
let mut previous_fallback = 0u64;
|
||||||
|
let mut previous_rejected = 0u64;
|
||||||
|
let pool_trim_low = buffer_pool
|
||||||
|
.max_buffers()
|
||||||
|
.min(BUFFER_POOL_TRIM_LOW_WATERMARK);
|
||||||
|
let pool_trim_high = buffer_pool
|
||||||
|
.max_buffers()
|
||||||
|
.min(BUFFER_POOL_TRIM_HIGH_WATERMARK);
|
||||||
|
let mut pool_trim_armed = true;
|
||||||
|
|
||||||
|
loop {
|
||||||
|
interval.tick().await;
|
||||||
|
let sample = read_system_memory_sample().await;
|
||||||
|
budget.update_system_sample(sample);
|
||||||
|
|
||||||
|
let snapshot = budget.snapshot();
|
||||||
|
let denied_delta = snapshot
|
||||||
|
.promotion_denied_total
|
||||||
|
.saturating_sub(previous_denied);
|
||||||
|
previous_denied = snapshot.promotion_denied_total;
|
||||||
|
let fallback_delta = snapshot
|
||||||
|
.minimum_fallback_total
|
||||||
|
.saturating_sub(previous_fallback);
|
||||||
|
previous_fallback = snapshot.minimum_fallback_total;
|
||||||
|
let rejected_delta = snapshot
|
||||||
|
.admission_rejected_total
|
||||||
|
.saturating_sub(previous_rejected);
|
||||||
|
previous_rejected = snapshot.admission_rejected_total;
|
||||||
|
|
||||||
|
let connection_pct = connection_fill_pct(stats.as_ref(), max_connections);
|
||||||
|
let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes);
|
||||||
|
let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes);
|
||||||
|
let pressure = shared.conntrack_pressure_active()
|
||||||
|
|| connection_pct.is_some_and(|value| value >= 85)
|
||||||
|
|| memory_available_pct.is_some_and(|value| value <= 15)
|
||||||
|
|| target_utilization_pct.is_some_and(|value| value >= 90)
|
||||||
|
|| denied_delta > 0
|
||||||
|
|| fallback_delta > 0
|
||||||
|
|| rejected_delta > 0;
|
||||||
|
|
||||||
|
if !pressure {
|
||||||
|
pool_trim_armed = true;
|
||||||
|
} else if pool_trim_armed && buffer_pool.pooled() > pool_trim_high {
|
||||||
|
buffer_pool.trim_to(pool_trim_low);
|
||||||
|
pool_trim_armed = false;
|
||||||
|
}
|
||||||
|
|
||||||
|
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),
|
||||||
|
);
|
||||||
|
stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard);
|
||||||
|
|
||||||
|
let headroom_target = if sample.total_bytes == 0 {
|
||||||
|
snapshot.hard_limit_bytes
|
||||||
|
} else {
|
||||||
|
snapshot
|
||||||
|
.reserved_bytes
|
||||||
|
.saturating_add(sample.available_bytes / 4)
|
||||||
|
.min(snapshot.hard_limit_bytes)
|
||||||
|
};
|
||||||
|
|
||||||
|
if pressure {
|
||||||
|
healthy_streak = 0;
|
||||||
|
let reduced = snapshot.target_bytes.saturating_mul(3) / 4;
|
||||||
|
budget.set_target_bytes(reduced.min(headroom_target));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let healthy = memory_available_pct.is_none_or(|value| value >= 30)
|
||||||
|
&& connection_pct.is_none_or(|value| value <= 70);
|
||||||
|
if !healthy {
|
||||||
|
healthy_streak = 0;
|
||||||
|
if headroom_target < snapshot.target_bytes {
|
||||||
|
budget.set_target_bytes(headroom_target);
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
healthy_streak = healthy_streak.saturating_add(1);
|
||||||
|
if healthy_streak >= HEALTHY_RECOVERY_SAMPLES {
|
||||||
|
healthy_streak = 0;
|
||||||
|
let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024);
|
||||||
|
budget.set_target_bytes(
|
||||||
|
snapshot
|
||||||
|
.target_bytes
|
||||||
|
.saturating_add(increment)
|
||||||
|
.min(headroom_target),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn connection_fill_pct(stats: &Stats, max_connections: u32) -> Option<u8> {
|
||||||
|
if max_connections == 0 {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(
|
||||||
|
((stats.get_current_connections_total().saturating_mul(100)) / u64::from(max_connections))
|
||||||
|
.min(100) as u8,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn percentage(value: u64, total: u64) -> Option<u8> {
|
||||||
|
if total == 0 {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(((value.saturating_mul(100)) / total).min(100) as u8)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn read_system_memory_sample() -> SystemMemorySample {
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
{
|
||||||
|
let meminfo = tokio::fs::read_to_string("/proc/meminfo")
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
let status = tokio::fs::read_to_string("/proc/self/status")
|
||||||
|
.await
|
||||||
|
.unwrap_or_default();
|
||||||
|
let host_total = parse_kib_field(&meminfo, "MemTotal:");
|
||||||
|
let host_available = parse_kib_field(&meminfo, "MemAvailable:");
|
||||||
|
let process_rss = parse_kib_field(&status, "VmRSS:");
|
||||||
|
|
||||||
|
let cgroup_v2_max = read_cgroup_limit("/sys/fs/cgroup/memory.max").await;
|
||||||
|
let cgroup_v2_current = read_u64_file("/sys/fs/cgroup/memory.current").await;
|
||||||
|
let cgroup_v1_max = read_cgroup_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes").await;
|
||||||
|
let cgroup_v1_current = read_u64_file("/sys/fs/cgroup/memory/memory.usage_in_bytes").await;
|
||||||
|
let cgroup_max = cgroup_v2_max.or(cgroup_v1_max);
|
||||||
|
let cgroup_current = cgroup_v2_current.or(cgroup_v1_current);
|
||||||
|
|
||||||
|
let total = match (host_total, cgroup_max) {
|
||||||
|
(0, Some(limit)) => limit,
|
||||||
|
(host, Some(limit)) => host.min(limit),
|
||||||
|
(host, None) => host,
|
||||||
|
};
|
||||||
|
let cgroup_available = cgroup_max
|
||||||
|
.zip(cgroup_current)
|
||||||
|
.map(|(limit, current)| limit.saturating_sub(current));
|
||||||
|
let available = match (host_available, cgroup_available) {
|
||||||
|
(0, Some(value)) => value,
|
||||||
|
(host, Some(value)) => host.min(value),
|
||||||
|
(host, None) => host,
|
||||||
|
};
|
||||||
|
return SystemMemorySample {
|
||||||
|
total_bytes: total,
|
||||||
|
available_bytes: available,
|
||||||
|
process_rss_bytes: process_rss,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
#[cfg(not(target_os = "linux"))]
|
||||||
|
{
|
||||||
|
SystemMemorySample::default()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
async fn read_cgroup_limit(path: &str) -> Option<u64> {
|
||||||
|
let raw = tokio::fs::read_to_string(path).await.ok()?;
|
||||||
|
let raw = raw.trim();
|
||||||
|
if raw == "max" {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let value = raw.parse::<u64>().ok()?;
|
||||||
|
(value < (1u64 << 60)).then_some(value)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
async fn read_u64_file(path: &str) -> Option<u64> {
|
||||||
|
tokio::fs::read_to_string(path)
|
||||||
|
.await
|
||||||
|
.ok()?
|
||||||
|
.trim()
|
||||||
|
.parse()
|
||||||
|
.ok()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(target_os = "linux")]
|
||||||
|
fn parse_kib_field(raw: &str, key: &str) -> u64 {
|
||||||
|
raw.lines()
|
||||||
|
.find_map(|line| {
|
||||||
|
let value = line.strip_prefix(key)?.split_whitespace().next()?;
|
||||||
|
value.parse::<u64>().ok()
|
||||||
|
})
|
||||||
|
.unwrap_or(0)
|
||||||
|
.saturating_mul(1024)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn align_up(bytes: usize) -> usize {
|
||||||
|
bytes
|
||||||
|
.div_ceil(DIRECT_BUFFER_UNIT_BYTES)
|
||||||
|
.saturating_mul(DIRECT_BUFFER_UNIT_BYTES)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn align_down(bytes: usize) -> usize {
|
||||||
|
bytes / DIRECT_BUFFER_UNIT_BYTES * DIRECT_BUFFER_UNIT_BYTES
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decrement_saturating(value: &AtomicU64) {
|
||||||
|
let mut current = value.load(Ordering::Relaxed);
|
||||||
|
while current != 0 {
|
||||||
|
match value.compare_exchange_weak(
|
||||||
|
current,
|
||||||
|
current - 1,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
) {
|
||||||
|
Ok(_) => return,
|
||||||
|
Err(observed) => current = observed,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
#[path = "tests/direct_buffer_budget_tests.rs"]
|
||||||
|
mod tests;
|
||||||
+16
-17
@@ -345,22 +345,22 @@ where
|
|||||||
} else {
|
} else {
|
||||||
Duration::from_secs(1800)
|
Duration::from_secs(1800)
|
||||||
};
|
};
|
||||||
let relay_result =
|
let relay_result = crate::proxy::relay::relay_direct_adaptive(
|
||||||
crate::proxy::relay::relay_bidirectional_with_activity_timeout_lease_and_cancel(
|
client_reader,
|
||||||
client_reader,
|
client_writer,
|
||||||
client_writer,
|
tg_reader,
|
||||||
tg_reader,
|
tg_writer,
|
||||||
tg_writer,
|
config.general.direct_relay_copy_buf_c2s_bytes,
|
||||||
config.general.direct_relay_copy_buf_c2s_bytes,
|
config.general.direct_relay_copy_buf_s2c_bytes,
|
||||||
config.general.direct_relay_copy_buf_s2c_bytes,
|
config.server.max_connections,
|
||||||
user,
|
user,
|
||||||
Arc::clone(&stats),
|
Arc::clone(&stats),
|
||||||
config.access.user_data_quota.get(user).copied(),
|
config.access.user_data_quota.get(user).copied(),
|
||||||
buffer_pool,
|
traffic_lease,
|
||||||
traffic_lease,
|
relay_activity_timeout,
|
||||||
relay_activity_timeout,
|
session_cancel.clone(),
|
||||||
session_cancel.clone(),
|
Arc::clone(&shared.direct_buffer_budget),
|
||||||
);
|
);
|
||||||
tokio::pin!(relay_result);
|
tokio::pin!(relay_result);
|
||||||
let relay_result = loop {
|
let relay_result = loop {
|
||||||
if let Some(cutover) =
|
if let Some(cutover) =
|
||||||
@@ -400,7 +400,6 @@ where
|
|||||||
Err(e) => debug!(user = %user, error = %e, "Direct relay ended with error"),
|
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();
|
let pool_snapshot = buffer_pool_trim.stats();
|
||||||
stats.set_buffer_pool_gauges(
|
stats.set_buffer_pool_gauges(
|
||||||
pool_snapshot.pooled,
|
pool_snapshot.pooled,
|
||||||
|
|||||||
@@ -149,6 +149,7 @@ where
|
|||||||
peer,
|
peer,
|
||||||
translated_local_addr,
|
translated_local_addr,
|
||||||
payload,
|
payload,
|
||||||
|
_permit,
|
||||||
flags,
|
flags,
|
||||||
effective_tag_array,
|
effective_tag_array,
|
||||||
)
|
)
|
||||||
@@ -822,7 +823,6 @@ where
|
|||||||
|
|
||||||
clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
|
clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
|
||||||
me_pool.registry().unregister(conn_id).await;
|
me_pool.registry().unregister(conn_id).await;
|
||||||
buffer_pool.trim_to(buffer_pool.max_buffers().min(64));
|
|
||||||
let pool_snapshot = buffer_pool.stats();
|
let pool_snapshot = buffer_pool.stats();
|
||||||
stats.set_buffer_pool_gauges(
|
stats.set_buffer_pool_gauges(
|
||||||
pool_snapshot.pooled,
|
pool_snapshot.pooled,
|
||||||
|
|||||||
@@ -60,6 +60,8 @@
|
|||||||
|
|
||||||
pub mod adaptive_buffers;
|
pub mod adaptive_buffers;
|
||||||
pub mod client;
|
pub mod client;
|
||||||
|
// Process-wide Direct relay copy-buffer ownership and pressure policy.
|
||||||
|
pub(crate) mod direct_buffer_budget;
|
||||||
pub mod direct_relay;
|
pub mod direct_relay;
|
||||||
pub mod handshake;
|
pub mod handshake;
|
||||||
pub mod masking;
|
pub mod masking;
|
||||||
|
|||||||
@@ -84,8 +84,11 @@ fn watchdog_delta(current: u64, previous: u64) -> u64 {
|
|||||||
current.saturating_sub(previous)
|
current.saturating_sub(previous)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
mod adaptive_copy;
|
||||||
mod io;
|
mod io;
|
||||||
|
|
||||||
|
pub(crate) use self::adaptive_copy::relay_direct_adaptive;
|
||||||
|
|
||||||
use self::io::{CombinedStream, SharedCounters, StatsIo, is_quota_io_error};
|
use self::io::{CombinedStream, SharedCounters, StatsIo, is_quota_io_error};
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
use self::io::{quota_adaptive_interval_bytes, should_immediate_quota_check};
|
use self::io::{quota_adaptive_interval_bytes, should_immediate_quota_check};
|
||||||
@@ -217,6 +220,7 @@ where
|
|||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[allow(dead_code)]
|
||||||
pub async fn relay_bidirectional_with_activity_timeout_lease_and_cancel<CR, CW, SR, SW>(
|
pub async fn relay_bidirectional_with_activity_timeout_lease_and_cancel<CR, CW, SR, SW>(
|
||||||
client_reader: CR,
|
client_reader: CR,
|
||||||
client_writer: CW,
|
client_writer: CW,
|
||||||
|
|||||||
@@ -0,0 +1,527 @@
|
|||||||
|
use std::io;
|
||||||
|
use std::pin::Pin;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf, copy_buf};
|
||||||
|
use tokio::time::Instant;
|
||||||
|
use tokio_util::sync::CancellationToken;
|
||||||
|
use tracing::{debug, warn};
|
||||||
|
|
||||||
|
use crate::error::{ProxyError, Result};
|
||||||
|
use crate::proxy::adaptive_buffers::{
|
||||||
|
AdaptiveTier, RelaySignalSample, SessionAdaptiveController, TierTransitionReason,
|
||||||
|
direct_copy_buffers_for_tier_with_ceilings,
|
||||||
|
};
|
||||||
|
use crate::proxy::direct_buffer_budget::{
|
||||||
|
DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES, DirectBufferBudget, DirectBufferLease,
|
||||||
|
};
|
||||||
|
use crate::proxy::traffic_limiter::TrafficLease;
|
||||||
|
use crate::stats::Stats;
|
||||||
|
|
||||||
|
use super::WATCHDOG_INTERVAL;
|
||||||
|
use super::io::{SharedCounters, StatsIo, is_quota_io_error};
|
||||||
|
use super::watchdog_delta;
|
||||||
|
|
||||||
|
mod write_pressure;
|
||||||
|
|
||||||
|
use self::write_pressure::WritePressureIo;
|
||||||
|
|
||||||
|
struct AdaptiveBufferState {
|
||||||
|
desired_bytes: AtomicUsize,
|
||||||
|
actual_bytes: AtomicUsize,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AdaptiveBufferState {
|
||||||
|
fn new(bytes: usize) -> Arc<Self> {
|
||||||
|
Arc::new(Self {
|
||||||
|
desired_bytes: AtomicUsize::new(bytes.max(1)),
|
||||||
|
actual_bytes: AtomicUsize::new(bytes.max(1)),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct AdaptiveBufReader<R> {
|
||||||
|
inner: R,
|
||||||
|
buffer: Box<[u8]>,
|
||||||
|
pos: usize,
|
||||||
|
cap: usize,
|
||||||
|
state: Arc<AdaptiveBufferState>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R> AdaptiveBufReader<R> {
|
||||||
|
fn new(inner: R, state: Arc<AdaptiveBufferState>) -> Self {
|
||||||
|
let bytes = state.actual_bytes.load(Ordering::Relaxed).max(1);
|
||||||
|
Self {
|
||||||
|
inner,
|
||||||
|
buffer: vec![0; bytes].into_boxed_slice(),
|
||||||
|
pos: 0,
|
||||||
|
cap: 0,
|
||||||
|
state,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resize_if_drained(&mut self) {
|
||||||
|
if self.pos != self.cap {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
let desired = self.state.desired_bytes.load(Ordering::Acquire).max(1);
|
||||||
|
if desired == self.buffer.len() {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
self.buffer = vec![0; desired].into_boxed_slice();
|
||||||
|
self.pos = 0;
|
||||||
|
self.cap = 0;
|
||||||
|
self.state.actual_bytes.store(desired, Ordering::Release);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: AsyncRead + Unpin> AsyncRead for AdaptiveBufReader<R> {
|
||||||
|
fn poll_read(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
cx: &mut Context<'_>,
|
||||||
|
output: &mut ReadBuf<'_>,
|
||||||
|
) -> Poll<io::Result<()>> {
|
||||||
|
let this = self.get_mut();
|
||||||
|
if this.pos < this.cap {
|
||||||
|
let available = &this.buffer[this.pos..this.cap];
|
||||||
|
let copied = available.len().min(output.remaining());
|
||||||
|
output.put_slice(&available[..copied]);
|
||||||
|
this.pos += copied;
|
||||||
|
return Poll::Ready(Ok(()));
|
||||||
|
}
|
||||||
|
this.resize_if_drained();
|
||||||
|
Pin::new(&mut this.inner).poll_read(cx, output)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<R: AsyncRead + Unpin> AsyncBufRead for AdaptiveBufReader<R> {
|
||||||
|
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
|
||||||
|
let this = self.get_mut();
|
||||||
|
if this.pos < this.cap {
|
||||||
|
return Poll::Ready(Ok(&this.buffer[this.pos..this.cap]));
|
||||||
|
}
|
||||||
|
|
||||||
|
this.resize_if_drained();
|
||||||
|
let mut read_buf = ReadBuf::new(&mut this.buffer);
|
||||||
|
match Pin::new(&mut this.inner).poll_read(cx, &mut read_buf) {
|
||||||
|
Poll::Ready(Ok(())) => {
|
||||||
|
this.pos = 0;
|
||||||
|
this.cap = read_buf.filled().len();
|
||||||
|
Poll::Ready(Ok(&this.buffer[..this.cap]))
|
||||||
|
}
|
||||||
|
Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
|
||||||
|
Poll::Pending => Poll::Pending,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn consume(self: Pin<&mut Self>, amount: usize) {
|
||||||
|
let this = self.get_mut();
|
||||||
|
this.pos = this.pos.saturating_add(amount).min(this.cap);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
enum AdaptiveRelayOutcome {
|
||||||
|
Copy(io::Result<(u64, u64)>),
|
||||||
|
ActivityTimeout,
|
||||||
|
UserDisabled,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
/// Relays one Direct session with independently resizable directional buffers.
|
||||||
|
pub(crate) async fn relay_direct_adaptive<CR, CW, SR, SW>(
|
||||||
|
client_reader: CR,
|
||||||
|
client_writer: CW,
|
||||||
|
server_reader: SR,
|
||||||
|
server_writer: SW,
|
||||||
|
ceiling_c2s_bytes: usize,
|
||||||
|
ceiling_s2c_bytes: usize,
|
||||||
|
max_connections: u32,
|
||||||
|
user: &str,
|
||||||
|
stats: Arc<Stats>,
|
||||||
|
quota_limit: Option<u64>,
|
||||||
|
traffic_lease: Option<Arc<TrafficLease>>,
|
||||||
|
activity_timeout: Duration,
|
||||||
|
session_cancel: CancellationToken,
|
||||||
|
budget: Arc<DirectBufferBudget>,
|
||||||
|
) -> Result<()>
|
||||||
|
where
|
||||||
|
CR: AsyncRead + Unpin + Send + 'static,
|
||||||
|
CW: AsyncWrite + Unpin + Send + 'static,
|
||||||
|
SR: AsyncRead + Unpin + Send + 'static,
|
||||||
|
SW: AsyncWrite + Unpin + Send + 'static,
|
||||||
|
{
|
||||||
|
let activity_timeout = activity_timeout.max(Duration::from_secs(1));
|
||||||
|
let epoch = Instant::now();
|
||||||
|
let counters = Arc::new(SharedCounters::new());
|
||||||
|
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
||||||
|
let user_owned = user.to_string();
|
||||||
|
|
||||||
|
let (base_c2s, base_s2c) = initial_base_sizes(
|
||||||
|
ceiling_c2s_bytes,
|
||||||
|
ceiling_s2c_bytes,
|
||||||
|
max_connections,
|
||||||
|
budget.target_bytes(),
|
||||||
|
);
|
||||||
|
let base_total = base_c2s.saturating_add(base_s2c);
|
||||||
|
let mut lease = match budget.try_reserve(base_total, false) {
|
||||||
|
Some(lease) => lease,
|
||||||
|
None => {
|
||||||
|
let minimum_total = DIRECT_BASE_C2S_BYTES + DIRECT_BASE_S2C_BYTES;
|
||||||
|
match budget.try_reserve(minimum_total, true) {
|
||||||
|
Some(lease) => {
|
||||||
|
budget.increment_minimum_fallback();
|
||||||
|
lease
|
||||||
|
}
|
||||||
|
None => {
|
||||||
|
budget.increment_admission_rejected();
|
||||||
|
return Err(ProxyError::Proxy(
|
||||||
|
"Direct relay buffer pressure: budget exhausted".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let effective_base = if lease.reserved_bytes() < base_total {
|
||||||
|
(DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES)
|
||||||
|
} else {
|
||||||
|
(base_c2s, base_s2c)
|
||||||
|
};
|
||||||
|
let c2s_state = AdaptiveBufferState::new(effective_base.0);
|
||||||
|
let s2c_state = AdaptiveBufferState::new(effective_base.1);
|
||||||
|
|
||||||
|
let mut controller = SessionAdaptiveController::new(AdaptiveTier::Base);
|
||||||
|
|
||||||
|
let c2s_client = StatsIo::new_with_traffic_lease(
|
||||||
|
client_reader,
|
||||||
|
Arc::clone(&counters),
|
||||||
|
Arc::clone(&stats),
|
||||||
|
user_owned.clone(),
|
||||||
|
traffic_lease.clone(),
|
||||||
|
quota_limit,
|
||||||
|
Arc::clone("a_exceeded),
|
||||||
|
epoch,
|
||||||
|
);
|
||||||
|
let client_writer = StatsIo::new_with_traffic_lease(
|
||||||
|
client_writer,
|
||||||
|
Arc::clone(&counters),
|
||||||
|
Arc::clone(&stats),
|
||||||
|
user_owned.clone(),
|
||||||
|
traffic_lease,
|
||||||
|
quota_limit,
|
||||||
|
Arc::clone("a_exceeded),
|
||||||
|
epoch,
|
||||||
|
);
|
||||||
|
let mut client_writer = WritePressureIo::new(client_writer, Arc::clone(&counters));
|
||||||
|
let mut c2s_reader = AdaptiveBufReader::new(c2s_client, Arc::clone(&c2s_state));
|
||||||
|
let mut s2c_reader = AdaptiveBufReader::new(server_reader, Arc::clone(&s2c_state));
|
||||||
|
let mut server_writer = server_writer;
|
||||||
|
let mut pressure_rx = budget.subscribe_pressure();
|
||||||
|
|
||||||
|
let relay_outcome = {
|
||||||
|
let copy = async {
|
||||||
|
let c2s = async {
|
||||||
|
let copied = copy_buf(&mut c2s_reader, &mut server_writer).await?;
|
||||||
|
server_writer.shutdown().await?;
|
||||||
|
Ok::<u64, io::Error>(copied)
|
||||||
|
};
|
||||||
|
let s2c = async {
|
||||||
|
let copied = copy_buf(&mut s2c_reader, &mut client_writer).await?;
|
||||||
|
client_writer.shutdown().await?;
|
||||||
|
Ok::<u64, io::Error>(copied)
|
||||||
|
};
|
||||||
|
tokio::try_join!(c2s, s2c)
|
||||||
|
};
|
||||||
|
tokio::pin!(copy);
|
||||||
|
|
||||||
|
let mut interval = tokio::time::interval(WATCHDOG_INTERVAL);
|
||||||
|
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||||
|
interval.tick().await;
|
||||||
|
let mut previous = RelaySignalSample::default();
|
||||||
|
let mut previous_log_c2s = 0u64;
|
||||||
|
let mut previous_log_s2c = 0u64;
|
||||||
|
let mut previous_sample_at = epoch;
|
||||||
|
|
||||||
|
loop {
|
||||||
|
tokio::select! {
|
||||||
|
result = &mut copy => break AdaptiveRelayOutcome::Copy(result),
|
||||||
|
_ = session_cancel.cancelled() => break AdaptiveRelayOutcome::UserDisabled,
|
||||||
|
changed = pressure_rx.changed() => {
|
||||||
|
if changed.is_ok() {
|
||||||
|
apply_global_pressure_demotion(
|
||||||
|
&mut controller,
|
||||||
|
&mut lease,
|
||||||
|
&c2s_state,
|
||||||
|
&s2c_state,
|
||||||
|
effective_base,
|
||||||
|
(ceiling_c2s_bytes, ceiling_s2c_bytes),
|
||||||
|
budget.as_ref(),
|
||||||
|
);
|
||||||
|
reconcile_reservation(&mut lease, &c2s_state, &s2c_state);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
_ = interval.tick() => {
|
||||||
|
let now = Instant::now();
|
||||||
|
let idle = counters.idle_duration(now, epoch);
|
||||||
|
if quota_exceeded.load(Ordering::Acquire) {
|
||||||
|
warn!(user = %user_owned, "User data quota reached, closing relay");
|
||||||
|
break AdaptiveRelayOutcome::ActivityTimeout;
|
||||||
|
}
|
||||||
|
if idle >= activity_timeout {
|
||||||
|
warn!(
|
||||||
|
user = %user_owned,
|
||||||
|
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||||
|
idle_secs = idle.as_secs(),
|
||||||
|
"Activity timeout"
|
||||||
|
);
|
||||||
|
break AdaptiveRelayOutcome::ActivityTimeout;
|
||||||
|
}
|
||||||
|
|
||||||
|
let sample = current_sample(counters.as_ref());
|
||||||
|
let c2s_delta = watchdog_delta(sample.c2s_bytes, previous_log_c2s);
|
||||||
|
let s2c_delta = watchdog_delta(sample.s2c_written_bytes, previous_log_s2c);
|
||||||
|
if c2s_delta > 0 || s2c_delta > 0 {
|
||||||
|
let secs = now.saturating_duration_since(previous_sample_at).as_secs_f64();
|
||||||
|
debug!(
|
||||||
|
user = %user_owned,
|
||||||
|
c2s_kbps = (c2s_delta as f64 / secs / 1024.0) as u64,
|
||||||
|
s2c_kbps = (s2c_delta as f64 / secs / 1024.0) as u64,
|
||||||
|
c2s_total = sample.c2s_bytes,
|
||||||
|
s2c_total = sample.s2c_written_bytes,
|
||||||
|
"Relay active"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let delta = sample_delta(sample, previous);
|
||||||
|
let tick_secs = now.saturating_duration_since(previous_sample_at).as_secs_f64();
|
||||||
|
if let Some(transition) = controller.observe(delta, tick_secs) {
|
||||||
|
apply_controller_transition(
|
||||||
|
transition,
|
||||||
|
&mut controller,
|
||||||
|
&mut lease,
|
||||||
|
&c2s_state,
|
||||||
|
&s2c_state,
|
||||||
|
effective_base,
|
||||||
|
(ceiling_c2s_bytes, ceiling_s2c_bytes),
|
||||||
|
budget.as_ref(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
reconcile_reservation(&mut lease, &c2s_state, &s2c_state);
|
||||||
|
previous = sample;
|
||||||
|
previous_log_c2s = sample.c2s_bytes;
|
||||||
|
previous_log_s2c = sample.s2c_written_bytes;
|
||||||
|
previous_sample_at = now;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let _ = client_writer.shutdown().await;
|
||||||
|
let _ = server_writer.shutdown().await;
|
||||||
|
let c2s_ops = counters.c2s_ops.load(Ordering::Relaxed);
|
||||||
|
let s2c_ops = counters.s2c_ops.load(Ordering::Relaxed);
|
||||||
|
let duration = epoch.elapsed();
|
||||||
|
match relay_outcome {
|
||||||
|
AdaptiveRelayOutcome::Copy(Ok((c2s, s2c))) => {
|
||||||
|
debug!(
|
||||||
|
user = %user_owned,
|
||||||
|
c2s_bytes = c2s,
|
||||||
|
s2c_bytes = s2c,
|
||||||
|
c2s_msgs = c2s_ops,
|
||||||
|
s2c_msgs = s2c_ops,
|
||||||
|
duration_secs = duration.as_secs(),
|
||||||
|
"Relay finished"
|
||||||
|
);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
AdaptiveRelayOutcome::Copy(Err(error)) if is_quota_io_error(&error) => {
|
||||||
|
warn!(
|
||||||
|
user = %user_owned,
|
||||||
|
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||||
|
c2s_msgs = c2s_ops,
|
||||||
|
s2c_msgs = s2c_ops,
|
||||||
|
duration_secs = duration.as_secs(),
|
||||||
|
"Data quota reached, closing relay"
|
||||||
|
);
|
||||||
|
Err(ProxyError::DataQuotaExceeded { user: user_owned })
|
||||||
|
}
|
||||||
|
AdaptiveRelayOutcome::Copy(Err(error)) => {
|
||||||
|
debug!(
|
||||||
|
user = %user_owned,
|
||||||
|
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||||
|
c2s_msgs = c2s_ops,
|
||||||
|
s2c_msgs = s2c_ops,
|
||||||
|
duration_secs = duration.as_secs(),
|
||||||
|
error = %error,
|
||||||
|
"Relay error"
|
||||||
|
);
|
||||||
|
Err(error.into())
|
||||||
|
}
|
||||||
|
AdaptiveRelayOutcome::ActivityTimeout => {
|
||||||
|
debug!(
|
||||||
|
user = %user_owned,
|
||||||
|
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||||
|
c2s_msgs = c2s_ops,
|
||||||
|
s2c_msgs = s2c_ops,
|
||||||
|
duration_secs = duration.as_secs(),
|
||||||
|
"Relay finished (activity timeout)"
|
||||||
|
);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
AdaptiveRelayOutcome::UserDisabled => {
|
||||||
|
debug!(
|
||||||
|
user = %user_owned,
|
||||||
|
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
|
||||||
|
c2s_msgs = c2s_ops,
|
||||||
|
s2c_msgs = s2c_ops,
|
||||||
|
duration_secs = duration.as_secs(),
|
||||||
|
"Relay finished (user disabled)"
|
||||||
|
);
|
||||||
|
Err(ProxyError::UserDisabled { user: user_owned })
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn initial_base_sizes(
|
||||||
|
ceiling_c2s: usize,
|
||||||
|
ceiling_s2c: usize,
|
||||||
|
max_connections: u32,
|
||||||
|
target_bytes: usize,
|
||||||
|
) -> (usize, usize) {
|
||||||
|
let configured_total = ceiling_c2s.saturating_add(ceiling_s2c);
|
||||||
|
let configured_worst_case = configured_total.saturating_mul(max_connections as usize);
|
||||||
|
if max_connections != 0 && configured_worst_case <= target_bytes {
|
||||||
|
return (ceiling_c2s, ceiling_s2c);
|
||||||
|
}
|
||||||
|
(
|
||||||
|
DIRECT_BASE_C2S_BYTES.min(ceiling_c2s),
|
||||||
|
DIRECT_BASE_S2C_BYTES.min(ceiling_s2c),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn current_sample(counters: &SharedCounters) -> RelaySignalSample {
|
||||||
|
RelaySignalSample {
|
||||||
|
c2s_bytes: counters.c2s_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_requested_bytes: counters.s2c_requested_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_written_bytes: counters.s2c_bytes.load(Ordering::Relaxed),
|
||||||
|
s2c_write_ops: counters.s2c_ops.load(Ordering::Relaxed),
|
||||||
|
s2c_partial_writes: counters.s2c_partial_writes.load(Ordering::Relaxed),
|
||||||
|
s2c_consecutive_pending_writes: counters
|
||||||
|
.s2c_consecutive_pending_writes
|
||||||
|
.load(Ordering::Relaxed),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn sample_delta(current: RelaySignalSample, previous: RelaySignalSample) -> RelaySignalSample {
|
||||||
|
RelaySignalSample {
|
||||||
|
c2s_bytes: current.c2s_bytes.saturating_sub(previous.c2s_bytes),
|
||||||
|
s2c_requested_bytes: current
|
||||||
|
.s2c_requested_bytes
|
||||||
|
.saturating_sub(previous.s2c_requested_bytes),
|
||||||
|
s2c_written_bytes: current
|
||||||
|
.s2c_written_bytes
|
||||||
|
.saturating_sub(previous.s2c_written_bytes),
|
||||||
|
s2c_write_ops: current.s2c_write_ops.saturating_sub(previous.s2c_write_ops),
|
||||||
|
s2c_partial_writes: current
|
||||||
|
.s2c_partial_writes
|
||||||
|
.saturating_sub(previous.s2c_partial_writes),
|
||||||
|
s2c_consecutive_pending_writes: current.s2c_consecutive_pending_writes,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[allow(clippy::too_many_arguments)]
|
||||||
|
fn apply_controller_transition(
|
||||||
|
transition: crate::proxy::adaptive_buffers::TierTransition,
|
||||||
|
controller: &mut SessionAdaptiveController,
|
||||||
|
lease: &mut DirectBufferLease,
|
||||||
|
c2s_state: &AdaptiveBufferState,
|
||||||
|
s2c_state: &AdaptiveBufferState,
|
||||||
|
base: (usize, usize),
|
||||||
|
ceilings: (usize, usize),
|
||||||
|
budget: &DirectBufferBudget,
|
||||||
|
) {
|
||||||
|
let sizes = direct_copy_buffers_for_tier_with_ceilings(
|
||||||
|
transition.to,
|
||||||
|
base.0,
|
||||||
|
base.1,
|
||||||
|
ceilings.0,
|
||||||
|
ceilings.1,
|
||||||
|
);
|
||||||
|
if transition.to > transition.from {
|
||||||
|
if !lease.try_grow_to(sizes.0.saturating_add(sizes.1)) {
|
||||||
|
*controller = SessionAdaptiveController::new(transition.from);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
match transition.reason {
|
||||||
|
TierTransitionReason::QuietDemotion => budget.increment_quiet_demotion(),
|
||||||
|
TierTransitionReason::SustainedWritePressure => {
|
||||||
|
budget.increment_write_pressure_demotion();
|
||||||
|
}
|
||||||
|
TierTransitionReason::SoftConfirmed | TierTransitionReason::HardPressure => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
set_desired_sizes(c2s_state, s2c_state, sizes);
|
||||||
|
lease.set_tier(transition.to.as_u8() as usize);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_global_pressure_demotion(
|
||||||
|
controller: &mut SessionAdaptiveController,
|
||||||
|
lease: &mut DirectBufferLease,
|
||||||
|
c2s_state: &AdaptiveBufferState,
|
||||||
|
s2c_state: &AdaptiveBufferState,
|
||||||
|
base: (usize, usize),
|
||||||
|
ceilings: (usize, usize),
|
||||||
|
budget: &DirectBufferBudget,
|
||||||
|
) {
|
||||||
|
let current = controller.tier();
|
||||||
|
let target = current.demote();
|
||||||
|
if target == current {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
*controller = SessionAdaptiveController::new(target);
|
||||||
|
let sizes =
|
||||||
|
direct_copy_buffers_for_tier_with_ceilings(target, base.0, base.1, ceilings.0, ceilings.1);
|
||||||
|
set_desired_sizes(c2s_state, s2c_state, sizes);
|
||||||
|
lease.set_tier(target.as_u8() as usize);
|
||||||
|
budget.increment_global_pressure_demotion();
|
||||||
|
}
|
||||||
|
|
||||||
|
fn set_desired_sizes(
|
||||||
|
c2s_state: &AdaptiveBufferState,
|
||||||
|
s2c_state: &AdaptiveBufferState,
|
||||||
|
sizes: (usize, usize),
|
||||||
|
) {
|
||||||
|
c2s_state
|
||||||
|
.desired_bytes
|
||||||
|
.store(sizes.0.max(1), Ordering::Release);
|
||||||
|
s2c_state
|
||||||
|
.desired_bytes
|
||||||
|
.store(sizes.1.max(1), Ordering::Release);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reconcile_reservation(
|
||||||
|
lease: &mut DirectBufferLease,
|
||||||
|
c2s_state: &AdaptiveBufferState,
|
||||||
|
s2c_state: &AdaptiveBufferState,
|
||||||
|
) {
|
||||||
|
// Promotion reserves the desired allocation before either reader grows.
|
||||||
|
// Demotion keeps the actual allocation covered until its buffered bytes drain.
|
||||||
|
let covered_c2s = c2s_state
|
||||||
|
.actual_bytes
|
||||||
|
.load(Ordering::Acquire)
|
||||||
|
.max(c2s_state.desired_bytes.load(Ordering::Acquire));
|
||||||
|
let covered_s2c = s2c_state
|
||||||
|
.actual_bytes
|
||||||
|
.load(Ordering::Acquire)
|
||||||
|
.max(s2c_state.desired_bytes.load(Ordering::Acquire));
|
||||||
|
lease.shrink_to(covered_c2s.saturating_add(covered_s2c));
|
||||||
|
}
|
||||||
@@ -0,0 +1,72 @@
|
|||||||
|
use std::io;
|
||||||
|
use std::pin::Pin;
|
||||||
|
use std::sync::Arc;
|
||||||
|
use std::sync::atomic::Ordering;
|
||||||
|
use std::task::{Context, Poll};
|
||||||
|
|
||||||
|
use tokio::io::AsyncWrite;
|
||||||
|
|
||||||
|
use super::super::io::SharedCounters;
|
||||||
|
|
||||||
|
/// Direct-only writer wrapper that exposes bounded backpressure signals.
|
||||||
|
pub(super) struct WritePressureIo<W> {
|
||||||
|
inner: W,
|
||||||
|
counters: Arc<SharedCounters>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<W> WritePressureIo<W> {
|
||||||
|
/// Wraps the client writer without changing its I/O or error contract.
|
||||||
|
pub(super) fn new(inner: W, counters: Arc<SharedCounters>) -> Self {
|
||||||
|
Self { inner, counters }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<W: AsyncWrite + Unpin> AsyncWrite for WritePressureIo<W> {
|
||||||
|
fn poll_write(
|
||||||
|
self: Pin<&mut Self>,
|
||||||
|
cx: &mut Context<'_>,
|
||||||
|
buffer: &[u8],
|
||||||
|
) -> Poll<io::Result<usize>> {
|
||||||
|
let this = self.get_mut();
|
||||||
|
if !buffer.is_empty() {
|
||||||
|
this.counters
|
||||||
|
.s2c_requested_bytes
|
||||||
|
.fetch_add(buffer.len() as u64, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
match Pin::new(&mut this.inner).poll_write(cx, buffer) {
|
||||||
|
Poll::Ready(Ok(written)) => {
|
||||||
|
this.counters
|
||||||
|
.s2c_consecutive_pending_writes
|
||||||
|
.store(0, Ordering::Relaxed);
|
||||||
|
if written < buffer.len() {
|
||||||
|
this.counters
|
||||||
|
.s2c_partial_writes
|
||||||
|
.fetch_add(1, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
Poll::Ready(Ok(written))
|
||||||
|
}
|
||||||
|
Poll::Ready(Err(error)) => {
|
||||||
|
this.counters
|
||||||
|
.s2c_consecutive_pending_writes
|
||||||
|
.store(0, Ordering::Relaxed);
|
||||||
|
Poll::Ready(Err(error))
|
||||||
|
}
|
||||||
|
Poll::Pending => {
|
||||||
|
let _ = this.counters.s2c_consecutive_pending_writes.fetch_update(
|
||||||
|
Ordering::Relaxed,
|
||||||
|
Ordering::Relaxed,
|
||||||
|
|current| Some(current.saturating_add(1)),
|
||||||
|
);
|
||||||
|
Poll::Pending
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||||
|
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
||||||
|
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,4 +1,4 @@
|
|||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use tokio::time::Instant;
|
use tokio::time::Instant;
|
||||||
@@ -20,6 +20,12 @@ pub(in crate::proxy::relay) struct SharedCounters {
|
|||||||
pub(in crate::proxy::relay) c2s_ops: AtomicU64,
|
pub(in crate::proxy::relay) c2s_ops: AtomicU64,
|
||||||
/// Number of poll_write completions (≈ S→C chunks)
|
/// Number of poll_write completions (≈ S→C chunks)
|
||||||
pub(in crate::proxy::relay) s2c_ops: AtomicU64,
|
pub(in crate::proxy::relay) s2c_ops: AtomicU64,
|
||||||
|
/// Bytes presented to client writes, including retried pending writes.
|
||||||
|
pub(in crate::proxy::relay) s2c_requested_bytes: AtomicU64,
|
||||||
|
/// Successful client writes that consumed only part of the offered slice.
|
||||||
|
pub(in crate::proxy::relay) s2c_partial_writes: AtomicU64,
|
||||||
|
/// Consecutive pending client writes observed by the active copy loop.
|
||||||
|
pub(in crate::proxy::relay) s2c_consecutive_pending_writes: AtomicU32,
|
||||||
/// Milliseconds since relay epoch of last I/O activity
|
/// Milliseconds since relay epoch of last I/O activity
|
||||||
last_activity_ms: AtomicU64,
|
last_activity_ms: AtomicU64,
|
||||||
}
|
}
|
||||||
@@ -31,6 +37,9 @@ impl SharedCounters {
|
|||||||
s2c_bytes: AtomicU64::new(0),
|
s2c_bytes: AtomicU64::new(0),
|
||||||
c2s_ops: AtomicU64::new(0),
|
c2s_ops: AtomicU64::new(0),
|
||||||
s2c_ops: AtomicU64::new(0),
|
s2c_ops: AtomicU64::new(0),
|
||||||
|
s2c_requested_bytes: AtomicU64::new(0),
|
||||||
|
s2c_partial_writes: AtomicU64::new(0),
|
||||||
|
s2c_consecutive_pending_writes: AtomicU32::new(0),
|
||||||
last_activity_ms: AtomicU64::new(0),
|
last_activity_ms: AtomicU64::new(0),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ use dashmap::DashMap;
|
|||||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
|
use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
|
||||||
use tokio_util::sync::CancellationToken;
|
use tokio_util::sync::CancellationToken;
|
||||||
|
|
||||||
|
use crate::proxy::direct_buffer_budget::{DirectBufferBudget, fallback_direct_buffer_hard_limit};
|
||||||
use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState};
|
use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState};
|
||||||
use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
|
use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
|
||||||
use crate::proxy::traffic_limiter::TrafficLimiter;
|
use crate::proxy::traffic_limiter::TrafficLimiter;
|
||||||
@@ -69,6 +70,7 @@ pub(crate) struct ProxySharedState {
|
|||||||
pub(crate) handshake: HandshakeSharedState,
|
pub(crate) handshake: HandshakeSharedState,
|
||||||
pub(crate) middle_relay: MiddleRelaySharedState,
|
pub(crate) middle_relay: MiddleRelaySharedState,
|
||||||
pub(crate) traffic_limiter: Arc<TrafficLimiter>,
|
pub(crate) traffic_limiter: Arc<TrafficLimiter>,
|
||||||
|
pub(crate) direct_buffer_budget: Arc<DirectBufferBudget>,
|
||||||
disabled_users: DashMap<String, ()>,
|
disabled_users: DashMap<String, ()>,
|
||||||
active_user_sessions: DashMap<(String, u64), CancellationToken>,
|
active_user_sessions: DashMap<(String, u64), CancellationToken>,
|
||||||
pub(crate) conntrack_pressure_active: AtomicBool,
|
pub(crate) conntrack_pressure_active: AtomicBool,
|
||||||
@@ -101,6 +103,15 @@ impl Drop for UserSessionGuard {
|
|||||||
|
|
||||||
impl ProxySharedState {
|
impl ProxySharedState {
|
||||||
pub(crate) fn new() -> Arc<Self> {
|
pub(crate) fn new() -> Arc<Self> {
|
||||||
|
Self::new_with_direct_buffer_budget(DirectBufferBudget::new(
|
||||||
|
fallback_direct_buffer_hard_limit(),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Creates process state with the startup-resolved Direct buffer envelope.
|
||||||
|
pub(crate) fn new_with_direct_buffer_budget(
|
||||||
|
direct_buffer_budget: Arc<DirectBufferBudget>,
|
||||||
|
) -> Arc<Self> {
|
||||||
Arc::new(Self {
|
Arc::new(Self {
|
||||||
handshake: HandshakeSharedState {
|
handshake: HandshakeSharedState {
|
||||||
auth_probe: DashMap::new(),
|
auth_probe: DashMap::new(),
|
||||||
@@ -129,6 +140,7 @@ impl ProxySharedState {
|
|||||||
relay_idle_mark_seq: AtomicU64::new(0),
|
relay_idle_mark_seq: AtomicU64::new(0),
|
||||||
},
|
},
|
||||||
traffic_limiter: TrafficLimiter::new(),
|
traffic_limiter: TrafficLimiter::new(),
|
||||||
|
direct_buffer_budget,
|
||||||
disabled_users: DashMap::new(),
|
disabled_users: DashMap::new(),
|
||||||
active_user_sessions: DashMap::new(),
|
active_user_sessions: DashMap::new(),
|
||||||
conntrack_pressure_active: AtomicBool::new(false),
|
conntrack_pressure_active: AtomicBool::new(false),
|
||||||
|
|||||||
@@ -0,0 +1,58 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn configured_direct_sizes_are_strict_tier_ceilings() {
|
||||||
|
let base = (4 * 1024, 8 * 1024);
|
||||||
|
let ceilings = (64 * 1024, 256 * 1024);
|
||||||
|
assert_eq!(
|
||||||
|
direct_copy_buffers_for_tier_with_ceilings(
|
||||||
|
AdaptiveTier::Base,
|
||||||
|
base.0,
|
||||||
|
base.1,
|
||||||
|
ceilings.0,
|
||||||
|
ceilings.1,
|
||||||
|
),
|
||||||
|
base
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
direct_copy_buffers_for_tier_with_ceilings(
|
||||||
|
AdaptiveTier::Tier3,
|
||||||
|
base.0,
|
||||||
|
base.1,
|
||||||
|
ceilings.0,
|
||||||
|
ceilings.1,
|
||||||
|
),
|
||||||
|
ceilings
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn sustained_pending_pressure_demotes_after_transient_promotion() {
|
||||||
|
let mut controller = SessionAdaptiveController::new(AdaptiveTier::Tier1);
|
||||||
|
let pressure = RelaySignalSample {
|
||||||
|
c2s_bytes: 0,
|
||||||
|
s2c_requested_bytes: 1024,
|
||||||
|
s2c_written_bytes: 0,
|
||||||
|
s2c_write_ops: 0,
|
||||||
|
s2c_partial_writes: 0,
|
||||||
|
s2c_consecutive_pending_writes: 3,
|
||||||
|
};
|
||||||
|
|
||||||
|
let first = controller
|
||||||
|
.observe(pressure, 10.0)
|
||||||
|
.expect("transient pressure must retain the staged promotion");
|
||||||
|
assert_eq!(first.reason, TierTransitionReason::HardPressure);
|
||||||
|
|
||||||
|
let second = controller
|
||||||
|
.observe(pressure, 10.0)
|
||||||
|
.expect("bounded transient pressure may promote one additional tier");
|
||||||
|
assert_eq!(second.reason, TierTransitionReason::HardPressure);
|
||||||
|
let sustained = controller
|
||||||
|
.observe(pressure, 10.0)
|
||||||
|
.expect("sustained pressure must release one tier");
|
||||||
|
assert_eq!(
|
||||||
|
sustained.reason,
|
||||||
|
TierTransitionReason::SustainedWritePressure
|
||||||
|
);
|
||||||
|
assert_eq!(sustained.to, AdaptiveTier::Tier2);
|
||||||
|
}
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn lease_drop_releases_the_complete_reservation() {
|
||||||
|
let budget = DirectBufferBudget::new(16 * 1024);
|
||||||
|
{
|
||||||
|
let lease = budget
|
||||||
|
.try_reserve(12 * 1024, false)
|
||||||
|
.expect("minimum reservation must fit");
|
||||||
|
assert_eq!(lease.reserved_bytes(), 12 * 1024);
|
||||||
|
assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024);
|
||||||
|
}
|
||||||
|
assert_eq!(budget.snapshot().reserved_bytes, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn absolute_ceiling_rejects_excess_minimum_reservations() {
|
||||||
|
let budget = DirectBufferBudget::new(16 * 1024);
|
||||||
|
let _lease = budget
|
||||||
|
.try_reserve(12 * 1024, true)
|
||||||
|
.expect("first minimum reservation must fit");
|
||||||
|
assert!(budget.try_reserve(8 * 1024, true).is_none());
|
||||||
|
assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn growth_and_shrink_keep_accounting_balanced() {
|
||||||
|
let budget = DirectBufferBudget::new(32 * 1024);
|
||||||
|
let mut lease = budget
|
||||||
|
.try_reserve(12 * 1024, false)
|
||||||
|
.expect("base reservation must fit");
|
||||||
|
assert!(lease.try_grow_to(24 * 1024));
|
||||||
|
assert_eq!(budget.snapshot().reserved_bytes, 24 * 1024);
|
||||||
|
lease.shrink_to(16 * 1024);
|
||||||
|
assert_eq!(budget.snapshot().reserved_bytes, 16 * 1024);
|
||||||
|
drop(lease);
|
||||||
|
assert_eq!(budget.snapshot().reserved_bytes, 0);
|
||||||
|
}
|
||||||
@@ -204,6 +204,12 @@ impl Stats {
|
|||||||
self.buffer_pool_in_use_gauge.load(Ordering::Relaxed)
|
self.buffer_pool_in_use_gauge.load(Ordering::Relaxed)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Returns the count of non-standard buffers replaced before pooling.
|
||||||
|
pub fn get_buffer_pool_replaced_nonstandard_total(&self) -> u64 {
|
||||||
|
self.buffer_pool_replaced_nonstandard_total
|
||||||
|
.load(Ordering::Relaxed)
|
||||||
|
}
|
||||||
|
|
||||||
pub fn get_me_c2me_send_full_total(&self) -> u64 {
|
pub fn get_me_c2me_send_full_total(&self) -> u64 {
|
||||||
self.me_c2me_send_full_total.load(Ordering::Relaxed)
|
self.me_c2me_send_full_total.load(Ordering::Relaxed)
|
||||||
}
|
}
|
||||||
@@ -269,6 +275,36 @@ 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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -274,6 +274,7 @@ pub struct Stats {
|
|||||||
buffer_pool_pooled_gauge: AtomicU64,
|
buffer_pool_pooled_gauge: AtomicU64,
|
||||||
buffer_pool_allocated_gauge: AtomicU64,
|
buffer_pool_allocated_gauge: AtomicU64,
|
||||||
buffer_pool_in_use_gauge: AtomicU64,
|
buffer_pool_in_use_gauge: AtomicU64,
|
||||||
|
buffer_pool_replaced_nonstandard_total: AtomicU64,
|
||||||
// C2ME enqueue observability
|
// C2ME enqueue observability
|
||||||
me_c2me_send_full_total: AtomicU64,
|
me_c2me_send_full_total: AtomicU64,
|
||||||
me_c2me_send_high_water_total: AtomicU64,
|
me_c2me_send_high_water_total: AtomicU64,
|
||||||
@@ -292,6 +293,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,
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -502,6 +552,14 @@ impl Stats {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Publishes the cumulative count of non-standard pool buffer replacements.
|
||||||
|
pub fn set_buffer_pool_replaced_nonstandard_total(&self, value: usize) {
|
||||||
|
if self.telemetry_me_allows_normal() {
|
||||||
|
self.buffer_pool_replaced_nonstandard_total
|
||||||
|
.store(value as u64, Ordering::Relaxed);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn increment_me_c2me_send_full_total(&self) {
|
pub fn increment_me_c2me_send_full_total(&self) {
|
||||||
if self.telemetry_me_allows_normal() {
|
if self.telemetry_me_allows_normal() {
|
||||||
self.me_c2me_send_full_total.fetch_add(1, Ordering::Relaxed);
|
self.me_c2me_send_full_total.fetch_add(1, Ordering::Relaxed);
|
||||||
|
|||||||
+20
-18
@@ -68,8 +68,8 @@ use crate::crypto::AesCtr;
|
|||||||
/// Actual limit is supplied at runtime from configuration.
|
/// Actual limit is supplied at runtime from configuration.
|
||||||
const DEFAULT_MAX_PENDING_WRITE: usize = 64 * 1024;
|
const DEFAULT_MAX_PENDING_WRITE: usize = 64 * 1024;
|
||||||
|
|
||||||
/// Default read buffer capacity (reader mostly decrypts in-place into caller buffer).
|
/// Maximum scratch capacity retained after a completed write.
|
||||||
const DEFAULT_READ_CAPACITY: usize = 16 * 1024;
|
const MAX_RETAINED_SCRATCH_CAPACITY: usize = 32 * 1024;
|
||||||
|
|
||||||
// ============= CryptoReader State =============
|
// ============= CryptoReader State =============
|
||||||
|
|
||||||
@@ -110,10 +110,6 @@ pub struct CryptoReader<R> {
|
|||||||
upstream: R,
|
upstream: R,
|
||||||
decryptor: AesCtr,
|
decryptor: AesCtr,
|
||||||
state: CryptoReaderState,
|
state: CryptoReaderState,
|
||||||
|
|
||||||
/// Reserved for future coalescing optimizations.
|
|
||||||
#[allow(dead_code)]
|
|
||||||
read_buf: BytesMut,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<R> CryptoReader<R> {
|
impl<R> CryptoReader<R> {
|
||||||
@@ -122,7 +118,6 @@ impl<R> CryptoReader<R> {
|
|||||||
upstream,
|
upstream,
|
||||||
decryptor,
|
decryptor,
|
||||||
state: CryptoReaderState::Idle,
|
state: CryptoReaderState::Idle,
|
||||||
read_buf: BytesMut::with_capacity(DEFAULT_READ_CAPACITY),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -321,7 +316,7 @@ struct PendingCiphertext {
|
|||||||
impl PendingCiphertext {
|
impl PendingCiphertext {
|
||||||
fn new(max_len: usize) -> Self {
|
fn new(max_len: usize) -> Self {
|
||||||
Self {
|
Self {
|
||||||
buf: BytesMut::with_capacity(16 * 1024),
|
buf: BytesMut::new(),
|
||||||
pos: 0,
|
pos: 0,
|
||||||
max_len,
|
max_len,
|
||||||
}
|
}
|
||||||
@@ -372,15 +367,12 @@ impl PendingCiphertext {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Replace the entire pending ciphertext by moving `src` in (swap, no copy).
|
/// Replace the entire pending ciphertext by moving `src` in without copying.
|
||||||
fn replace_with(&mut self, mut src: BytesMut) {
|
fn replace_with(&mut self, src: BytesMut) {
|
||||||
debug_assert!(src.len() <= self.max_len);
|
debug_assert!(src.len() <= self.max_len);
|
||||||
|
|
||||||
self.buf.clear();
|
self.buf = src;
|
||||||
self.pos = 0;
|
self.pos = 0;
|
||||||
|
|
||||||
// Swap: keep allocations hot and avoid copying bytes.
|
|
||||||
std::mem::swap(&mut self.buf, &mut src);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Append plaintext and encrypt appended range in-place.
|
/// Append plaintext and encrypt appended range in-place.
|
||||||
@@ -465,7 +457,7 @@ impl<W> CryptoWriter<W> {
|
|||||||
upstream,
|
upstream,
|
||||||
encryptor,
|
encryptor,
|
||||||
state: CryptoWriterState::Idle,
|
state: CryptoWriterState::Idle,
|
||||||
scratch: BytesMut::with_capacity(16 * 1024),
|
scratch: BytesMut::new(),
|
||||||
max_pending_write: max_pending.max(4 * 1024),
|
max_pending_write: max_pending.max(4 * 1024),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -502,6 +494,7 @@ impl<W> CryptoWriter<W> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn poison(&mut self, error: io::Error) {
|
fn poison(&mut self, error: io::Error) {
|
||||||
|
self.scratch = BytesMut::new();
|
||||||
self.state = CryptoWriterState::Poisoned { error: Some(error) };
|
self.state = CryptoWriterState::Poisoned { error: Some(error) };
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -552,6 +545,15 @@ impl<W> CryptoWriter<W> {
|
|||||||
scratch.extend_from_slice(plaintext);
|
scratch.extend_from_slice(plaintext);
|
||||||
encryptor.apply(&mut scratch[..]);
|
encryptor.apply(&mut scratch[..]);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Clear reusable scratch while releasing allocations inflated by large writes.
|
||||||
|
fn recycle_scratch(scratch: &mut BytesMut) {
|
||||||
|
if scratch.capacity() > MAX_RETAINED_SCRATCH_CAPACITY {
|
||||||
|
*scratch = BytesMut::new();
|
||||||
|
} else {
|
||||||
|
scratch.clear();
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<W: AsyncWrite + Unpin> CryptoWriter<W> {
|
impl<W: AsyncWrite + Unpin> CryptoWriter<W> {
|
||||||
@@ -698,13 +700,13 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for CryptoWriter<W> {
|
|||||||
|
|
||||||
Poll::Ready(Ok(n)) => {
|
Poll::Ready(Ok(n)) => {
|
||||||
if n == this.scratch.len() {
|
if n == this.scratch.len() {
|
||||||
this.scratch.clear();
|
Self::recycle_scratch(&mut this.scratch);
|
||||||
return Poll::Ready(Ok(to_accept));
|
return Poll::Ready(Ok(to_accept));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Partial upstream write of ciphertext
|
// Partial upstream write of ciphertext
|
||||||
let remainder = this.scratch.split_off(n);
|
let mut remainder = std::mem::take(&mut this.scratch);
|
||||||
this.scratch.clear();
|
let _ = remainder.split_to(n);
|
||||||
|
|
||||||
let pending = Self::ensure_pending(&mut this.state, this.max_pending_write);
|
let pending = Self::ensure_pending(&mut this.state, this.max_pending_write);
|
||||||
pending.replace_with(remainder);
|
pending.replace_with(remainder);
|
||||||
|
|||||||
+66
-34
@@ -40,7 +40,7 @@ use std::pin::Pin;
|
|||||||
use std::task::{Context, Poll};
|
use std::task::{Context, Poll};
|
||||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
|
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
|
||||||
|
|
||||||
use super::state::{HeaderBuffer, StreamState, WriteBuffer, YieldBuffer};
|
use super::state::{HeaderBuffer, StreamState, YieldBuffer};
|
||||||
use crate::protocol::constants::{
|
use crate::protocol::constants::{
|
||||||
MAX_TLS_CIPHERTEXT_SIZE, MAX_TLS_PLAINTEXT_SIZE, TLS_RECORD_ALERT, TLS_RECORD_APPLICATION,
|
MAX_TLS_CIPHERTEXT_SIZE, MAX_TLS_PLAINTEXT_SIZE, TLS_RECORD_ALERT, TLS_RECORD_APPLICATION,
|
||||||
TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE, TLS_VERSION,
|
TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE, TLS_VERSION,
|
||||||
@@ -59,6 +59,9 @@ const MAX_TLS_PAYLOAD: usize = MAX_TLS_CIPHERTEXT_SIZE;
|
|||||||
/// Note: we never queue unlimited amount of data here; state holds at most one record.
|
/// Note: we never queue unlimited amount of data here; state holds at most one record.
|
||||||
const MAX_PENDING_WRITE: usize = 64 * 1024;
|
const MAX_PENDING_WRITE: usize = 64 * 1024;
|
||||||
|
|
||||||
|
/// Maximum record buffer capacity retained between writes.
|
||||||
|
const MAX_RETAINED_RECORD_CAPACITY: usize = 2 * (TLS_HEADER_SIZE + MAX_TLS_PAYLOAD);
|
||||||
|
|
||||||
// ============= TLS Record Types =============
|
// ============= TLS Record Types =============
|
||||||
|
|
||||||
/// Parsed TLS record header (5 bytes)
|
/// Parsed TLS record header (5 bytes)
|
||||||
@@ -639,7 +642,7 @@ enum TlsWriterState {
|
|||||||
|
|
||||||
/// Writing a complete TLS record (header + body), possibly partially
|
/// Writing a complete TLS record (header + body), possibly partially
|
||||||
WritingRecord {
|
WritingRecord {
|
||||||
record: WriteBuffer,
|
position: usize,
|
||||||
payload_size: usize,
|
payload_size: usize,
|
||||||
},
|
},
|
||||||
|
|
||||||
@@ -676,6 +679,7 @@ impl StreamState for TlsWriterState {
|
|||||||
pub struct FakeTlsWriter<W> {
|
pub struct FakeTlsWriter<W> {
|
||||||
upstream: W,
|
upstream: W,
|
||||||
state: TlsWriterState,
|
state: TlsWriterState,
|
||||||
|
record_buffer: BytesMut,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<W> FakeTlsWriter<W> {
|
impl<W> FakeTlsWriter<W> {
|
||||||
@@ -683,6 +687,7 @@ impl<W> FakeTlsWriter<W> {
|
|||||||
Self {
|
Self {
|
||||||
upstream,
|
upstream,
|
||||||
state: TlsWriterState::Idle,
|
state: TlsWriterState::Idle,
|
||||||
|
record_buffer: BytesMut::new(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -704,10 +709,15 @@ impl<W> FakeTlsWriter<W> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn has_pending(&self) -> bool {
|
pub fn has_pending(&self) -> bool {
|
||||||
matches!(&self.state, TlsWriterState::WritingRecord { record, .. } if !record.is_empty())
|
matches!(
|
||||||
|
&self.state,
|
||||||
|
TlsWriterState::WritingRecord { position, .. }
|
||||||
|
if *position < self.record_buffer.len()
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn poison(&mut self, error: io::Error) {
|
fn poison(&mut self, error: io::Error) {
|
||||||
|
self.record_buffer = BytesMut::new();
|
||||||
self.state = TlsWriterState::Poisoned { error: Some(error) };
|
self.state = TlsWriterState::Poisoned { error: Some(error) };
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -720,17 +730,26 @@ impl<W> FakeTlsWriter<W> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn build_record(data: &[u8]) -> BytesMut {
|
fn build_record(&mut self, data: &[u8]) {
|
||||||
let header = TlsRecordHeader {
|
let header = TlsRecordHeader {
|
||||||
record_type: TLS_RECORD_APPLICATION,
|
record_type: TLS_RECORD_APPLICATION,
|
||||||
version: TLS_VERSION,
|
version: TLS_VERSION,
|
||||||
length: data.len() as u16,
|
length: data.len() as u16,
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut record = BytesMut::with_capacity(TLS_HEADER_SIZE + data.len());
|
self.record_buffer.clear();
|
||||||
record.extend_from_slice(&header.to_bytes());
|
self.record_buffer.reserve(TLS_HEADER_SIZE + data.len());
|
||||||
record.extend_from_slice(data);
|
self.record_buffer.extend_from_slice(&header.to_bytes());
|
||||||
record
|
self.record_buffer.extend_from_slice(data);
|
||||||
|
debug_assert!(self.record_buffer.len() <= MAX_PENDING_WRITE);
|
||||||
|
}
|
||||||
|
|
||||||
|
fn recycle_record_buffer(&mut self) {
|
||||||
|
if self.record_buffer.capacity() > MAX_RETAINED_RECORD_CAPACITY {
|
||||||
|
self.record_buffer = BytesMut::new();
|
||||||
|
} else {
|
||||||
|
self.record_buffer.clear();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -744,10 +763,11 @@ impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
|
|||||||
fn poll_flush_record_inner(
|
fn poll_flush_record_inner(
|
||||||
upstream: &mut W,
|
upstream: &mut W,
|
||||||
cx: &mut Context<'_>,
|
cx: &mut Context<'_>,
|
||||||
record: &mut WriteBuffer,
|
record: &[u8],
|
||||||
|
position: &mut usize,
|
||||||
) -> FlushResult {
|
) -> FlushResult {
|
||||||
while !record.is_empty() {
|
while *position < record.len() {
|
||||||
let data = record.pending();
|
let data = &record[*position..];
|
||||||
match Pin::new(&mut *upstream).poll_write(cx, data) {
|
match Pin::new(&mut *upstream).poll_write(cx, data) {
|
||||||
Poll::Pending => return FlushResult::Pending,
|
Poll::Pending => return FlushResult::Pending,
|
||||||
Poll::Ready(Err(e)) => return FlushResult::Error(e),
|
Poll::Ready(Err(e)) => return FlushResult::Error(e),
|
||||||
@@ -757,7 +777,7 @@ impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
|
|||||||
"upstream returned 0 bytes written",
|
"upstream returned 0 bytes written",
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
Poll::Ready(Ok(n)) => record.advance(n),
|
Poll::Ready(Ok(n)) => *position += n,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -780,14 +800,19 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
TlsWriterState::WritingRecord {
|
TlsWriterState::WritingRecord {
|
||||||
mut record,
|
mut position,
|
||||||
payload_size,
|
payload_size,
|
||||||
} => {
|
} => {
|
||||||
// Finish writing previous record before accepting new bytes.
|
// Finish writing previous record before accepting new bytes.
|
||||||
match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
|
match Self::poll_flush_record_inner(
|
||||||
|
&mut this.upstream,
|
||||||
|
cx,
|
||||||
|
&this.record_buffer,
|
||||||
|
&mut position,
|
||||||
|
) {
|
||||||
FlushResult::Pending => {
|
FlushResult::Pending => {
|
||||||
this.state = TlsWriterState::WritingRecord {
|
this.state = TlsWriterState::WritingRecord {
|
||||||
record,
|
position,
|
||||||
payload_size,
|
payload_size,
|
||||||
};
|
};
|
||||||
return Poll::Pending;
|
return Poll::Pending;
|
||||||
@@ -797,6 +822,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
|||||||
return Poll::Ready(Err(e));
|
return Poll::Ready(Err(e));
|
||||||
}
|
}
|
||||||
FlushResult::Complete(_) => {
|
FlushResult::Complete(_) => {
|
||||||
|
this.recycle_record_buffer();
|
||||||
this.state = TlsWriterState::Idle;
|
this.state = TlsWriterState::Idle;
|
||||||
// continue to accept new buf below
|
// continue to accept new buf below
|
||||||
}
|
}
|
||||||
@@ -818,19 +844,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
|||||||
let chunk = &buf[..chunk_size];
|
let chunk = &buf[..chunk_size];
|
||||||
|
|
||||||
// Build the complete record (header + payload)
|
// Build the complete record (header + payload)
|
||||||
let record_data = Self::build_record(chunk);
|
this.build_record(chunk);
|
||||||
|
|
||||||
match Pin::new(&mut this.upstream).poll_write(cx, &record_data) {
|
match Pin::new(&mut this.upstream).poll_write(cx, &this.record_buffer) {
|
||||||
Poll::Ready(Ok(n)) if n == record_data.len() => Poll::Ready(Ok(chunk_size)),
|
Poll::Ready(Ok(n)) if n == this.record_buffer.len() => {
|
||||||
|
this.recycle_record_buffer();
|
||||||
|
Poll::Ready(Ok(chunk_size))
|
||||||
|
}
|
||||||
|
|
||||||
Poll::Ready(Ok(n)) => {
|
Poll::Ready(Ok(n)) => {
|
||||||
// Partial write of the record: store remainder.
|
|
||||||
let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
|
|
||||||
// record_data length is <= 16389, fits MAX_PENDING_WRITE
|
|
||||||
let _ = write_buffer.extend(&record_data[n..]);
|
|
||||||
|
|
||||||
this.state = TlsWriterState::WritingRecord {
|
this.state = TlsWriterState::WritingRecord {
|
||||||
record: write_buffer,
|
position: n,
|
||||||
payload_size: chunk_size,
|
payload_size: chunk_size,
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -844,12 +868,8 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
Poll::Pending => {
|
Poll::Pending => {
|
||||||
// Buffer entire record and report success for this chunk.
|
|
||||||
let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
|
|
||||||
let _ = write_buffer.extend(&record_data);
|
|
||||||
|
|
||||||
this.state = TlsWriterState::WritingRecord {
|
this.state = TlsWriterState::WritingRecord {
|
||||||
record: write_buffer,
|
position: 0,
|
||||||
payload_size: chunk_size,
|
payload_size: chunk_size,
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -871,12 +891,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
TlsWriterState::WritingRecord {
|
TlsWriterState::WritingRecord {
|
||||||
mut record,
|
mut position,
|
||||||
payload_size,
|
payload_size,
|
||||||
} => match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
|
} => match Self::poll_flush_record_inner(
|
||||||
|
&mut this.upstream,
|
||||||
|
cx,
|
||||||
|
&this.record_buffer,
|
||||||
|
&mut position,
|
||||||
|
) {
|
||||||
FlushResult::Pending => {
|
FlushResult::Pending => {
|
||||||
this.state = TlsWriterState::WritingRecord {
|
this.state = TlsWriterState::WritingRecord {
|
||||||
record,
|
position,
|
||||||
payload_size,
|
payload_size,
|
||||||
};
|
};
|
||||||
return Poll::Pending;
|
return Poll::Pending;
|
||||||
@@ -886,6 +911,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
|||||||
return Poll::Ready(Err(e));
|
return Poll::Ready(Err(e));
|
||||||
}
|
}
|
||||||
FlushResult::Complete(_) => {
|
FlushResult::Complete(_) => {
|
||||||
|
this.recycle_record_buffer();
|
||||||
this.state = TlsWriterState::Idle;
|
this.state = TlsWriterState::Idle;
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
@@ -905,11 +931,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
|
|||||||
|
|
||||||
match state {
|
match state {
|
||||||
TlsWriterState::WritingRecord {
|
TlsWriterState::WritingRecord {
|
||||||
mut record,
|
mut position,
|
||||||
payload_size: _,
|
payload_size: _,
|
||||||
} => {
|
} => {
|
||||||
// Best-effort flush (do not block shutdown forever).
|
// Best-effort flush (do not block shutdown forever).
|
||||||
let _ = Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record);
|
let _ = Self::poll_flush_record_inner(
|
||||||
|
&mut this.upstream,
|
||||||
|
cx,
|
||||||
|
&this.record_buffer,
|
||||||
|
&mut position,
|
||||||
|
);
|
||||||
|
this.recycle_record_buffer();
|
||||||
this.state = TlsWriterState::Idle;
|
this.state = TlsWriterState::Idle;
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
|
|||||||
@@ -1,9 +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 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};
|
||||||
@@ -11,9 +14,66 @@ 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(Bytes),
|
Data {
|
||||||
|
payload: Bytes,
|
||||||
|
_permit: Option<OwnedSemaphorePermit>,
|
||||||
|
writer_permit: WriterBytePermit,
|
||||||
|
},
|
||||||
DataAndFlush(Bytes),
|
DataAndFlush(Bytes),
|
||||||
ProxyReq(ProxyReqCommand),
|
ProxyReq(ProxyReqCommand),
|
||||||
ControlAndFlush([u8; 12]),
|
ControlAndFlush([u8; 12]),
|
||||||
@@ -28,6 +88,8 @@ pub(crate) struct ProxyReqCommand {
|
|||||||
pub(crate) proto_flags: u32,
|
pub(crate) proto_flags: u32,
|
||||||
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) writer_permit: WriterBytePermit,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
@@ -82,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());
|
||||||
@@ -274,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());
|
||||||
@@ -322,3 +384,38 @@ 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)) => {
|
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();
|
||||||
|
|||||||
@@ -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)]
|
||||||
|
|||||||
@@ -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!(
|
||||||
|
|||||||
@@ -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,9 @@ 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 +426,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 +448,8 @@ impl ConnRegistry {
|
|||||||
Some((
|
Some((
|
||||||
ConnWriter {
|
ConnWriter {
|
||||||
writer_id,
|
writer_id,
|
||||||
tx: writer,
|
tx: writer.tx,
|
||||||
|
byte_budget: writer.byte_budget,
|
||||||
},
|
},
|
||||||
meta,
|
meta,
|
||||||
))
|
))
|
||||||
|
|||||||
@@ -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::mpsc;
|
|
||||||
use tokio::sync::mpsc::error::TrySendError;
|
use tokio::sync::mpsc::error::TrySendError;
|
||||||
|
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc};
|
||||||
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,34 +43,133 @@ 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())
|
||||||
}
|
}
|
||||||
|
|
||||||
fn proxy_req_payload_from_command(cmd: WriterCommand) -> Option<PooledBuffer> {
|
fn proxy_req_payload_from_command(
|
||||||
|
cmd: WriterCommand,
|
||||||
|
) -> Option<(PooledBuffer, OwnedSemaphorePermit)> {
|
||||||
match cmd {
|
match cmd {
|
||||||
WriterCommand::ProxyReq(command) => Some(command.payload),
|
WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn payload_permit_from_data_command(cmd: WriterCommand) -> Option<OwnedSemaphorePermit> {
|
||||||
|
match cmd {
|
||||||
|
WriterCommand::Data { _permit, .. } => _permit,
|
||||||
_ => None,
|
_ => None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
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) => {
|
||||||
Ok(Ok(permit)) => Ok(permit),
|
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve)
|
||||||
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
|
.await
|
||||||
Err(_) => Err(WriterCommandReserveError::TimedOut),
|
{
|
||||||
},
|
Ok(Ok(permit)) => Ok(permit),
|
||||||
|
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
|
||||||
|
Err(_) => Err(WriterCommandReserveError::TimedOut),
|
||||||
|
}
|
||||||
|
}
|
||||||
None => reserve.await.map_err(|_| WriterCommandReserveError::Closed),
|
None => reserve.await.map_err(|_| WriterCommandReserveError::Closed),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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.
|
||||||
pub async fn send_proxy_req(
|
pub async fn send_proxy_req(
|
||||||
self: &Arc<Self>,
|
self: &Arc<Self>,
|
||||||
conn_id: u64,
|
conn_id: u64,
|
||||||
@@ -76,8 +179,32 @@ impl MePool {
|
|||||||
data: &[u8],
|
data: &[u8],
|
||||||
proto_flags: u32,
|
proto_flags: u32,
|
||||||
tag_override: Option<&[u8]>,
|
tag_override: Option<&[u8]>,
|
||||||
|
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(
|
||||||
@@ -118,19 +245,48 @@ 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(
|
||||||
|
¤t.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);
|
||||||
match current.tx.try_send(WriterCommand::Data(current_payload)) {
|
let command = WriterCommand::Data {
|
||||||
|
payload: current_payload,
|
||||||
|
_permit: payload_permit.take(),
|
||||||
|
writer_permit,
|
||||||
|
};
|
||||||
|
match current.tx.try_send(command) {
|
||||||
Ok(()) => {
|
Ok(()) => {
|
||||||
self.note_hybrid_route_success();
|
self.note_hybrid_route_success();
|
||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
Err(TrySendError::Full(cmd)) => {
|
Err(TrySendError::Full(cmd)) => {
|
||||||
match reserve_writer_command_slot(
|
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||||
¤t.tx,
|
|
||||||
self.route_runtime.me_route_blocking_send_timeout,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(permit) => {
|
Ok(permit) => {
|
||||||
permit.send(cmd);
|
permit.send(cmd);
|
||||||
self.note_hybrid_route_success();
|
self.note_hybrid_route_success();
|
||||||
@@ -143,14 +299,17 @@ impl MePool {
|
|||||||
"ME writer channel full within blocking send timeout".into(),
|
"ME writer channel full within blocking send timeout".into(),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
Err(WriterCommandReserveError::Closed) => {}
|
Err(WriterCommandReserveError::Closed) => {
|
||||||
|
payload_permit = payload_permit_from_data_command(cmd);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||||
self.remove_writer_and_close_clients(current.writer_id)
|
self.remove_writer_and_close_clients(current.writer_id)
|
||||||
.await;
|
.await;
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
Err(TrySendError::Closed(_)) => {
|
Err(TrySendError::Closed(cmd)) => {
|
||||||
|
payload_permit = payload_permit_from_data_command(cmd);
|
||||||
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
warn!(writer_id = current.writer_id, "ME writer channel closed");
|
||||||
self.remove_writer_and_close_clients(current.writer_id)
|
self.remove_writer_and_close_clients(current.writer_id)
|
||||||
.await;
|
.await;
|
||||||
@@ -464,11 +623,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,
|
||||||
@@ -479,7 +658,11 @@ impl MePool {
|
|||||||
self.remove_writer_and_close_clients(w.id).await;
|
self.remove_writer_and_close_clients(w.id).await;
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
permit.send(WriterCommand::Data(payload));
|
permit.send(WriterCommand::Data {
|
||||||
|
payload,
|
||||||
|
_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);
|
||||||
if w.generation < self.current_generation() {
|
if w.generation < self.current_generation() {
|
||||||
@@ -521,52 +704,71 @@ 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(self.route_runtime.me_route_blocking_send_timeout);
|
||||||
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
|
let writer_permit = match reserve_writer_bytes(
|
||||||
let (payload, meta) = build_routed_payload(effective_our_addr);
|
&w.byte_budget,
|
||||||
let reserve_result =
|
writer_byte_permits,
|
||||||
if let Some(timeout) = self.route_runtime.me_route_blocking_send_timeout {
|
writer_reserved_bytes,
|
||||||
match tokio::time::timeout(timeout, w.tx.clone().reserve_owned()).await {
|
deadline,
|
||||||
Ok(result) => result,
|
&self.stats,
|
||||||
Err(_) => {
|
)
|
||||||
self.stats.increment_me_writer_pick_full_total(pick_mode);
|
.await
|
||||||
continue;
|
{
|
||||||
}
|
Ok(permit) => permit,
|
||||||
}
|
Err(WriterByteReserveError::TimedOut) => {
|
||||||
} else {
|
self.stats.increment_me_writer_pick_full_total(pick_mode);
|
||||||
w.tx.clone().reserve_owned().await
|
continue;
|
||||||
};
|
|
||||||
match reserve_result {
|
|
||||||
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));
|
|
||||||
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(());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light.
|
/// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light.
|
||||||
|
/// The client byte permit follows the payload until writer completion or command drop.
|
||||||
pub async fn send_proxy_req_pooled(
|
pub async fn send_proxy_req_pooled(
|
||||||
self: &Arc<Self>,
|
self: &Arc<Self>,
|
||||||
conn_id: u64,
|
conn_id: u64,
|
||||||
@@ -574,12 +776,69 @@ impl MePool {
|
|||||||
client_addr: SocketAddr,
|
client_addr: SocketAddr,
|
||||||
our_addr: SocketAddr,
|
our_addr: SocketAddr,
|
||||||
payload: PooledBuffer,
|
payload: PooledBuffer,
|
||||||
|
_permit: OwnedSemaphorePermit,
|
||||||
proto_flags: u32,
|
proto_flags: u32,
|
||||||
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(
|
||||||
|
¤t.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,
|
||||||
@@ -587,6 +846,8 @@ impl MePool {
|
|||||||
proto_flags,
|
proto_flags,
|
||||||
proxy_tag: tag,
|
proxy_tag: tag,
|
||||||
payload,
|
payload,
|
||||||
|
_permit,
|
||||||
|
writer_permit,
|
||||||
});
|
});
|
||||||
match current.tx.try_send(command) {
|
match current.tx.try_send(command) {
|
||||||
Ok(()) => {
|
Ok(()) => {
|
||||||
@@ -594,12 +855,7 @@ impl MePool {
|
|||||||
return Ok(());
|
return Ok(());
|
||||||
}
|
}
|
||||||
Err(TrySendError::Full(cmd)) => {
|
Err(TrySendError::Full(cmd)) => {
|
||||||
match reserve_writer_command_slot(
|
match reserve_writer_command_slot(¤t.tx, deadline).await {
|
||||||
¤t.tx,
|
|
||||||
self.route_runtime.me_route_blocking_send_timeout,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(permit) => {
|
Ok(permit) => {
|
||||||
permit.send(cmd);
|
permit.send(cmd);
|
||||||
self.note_hybrid_route_success();
|
self.note_hybrid_route_success();
|
||||||
@@ -613,7 +869,8 @@ impl MePool {
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
Err(WriterCommandReserveError::Closed) => {
|
Err(WriterCommandReserveError::Closed) => {
|
||||||
let Some(payload) = proxy_req_payload_from_command(cmd) else {
|
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd)
|
||||||
|
else {
|
||||||
return Err(ProxyError::Proxy(
|
return Err(ProxyError::Proxy(
|
||||||
"ME writer rejected unexpected command type".into(),
|
"ME writer rejected unexpected command type".into(),
|
||||||
));
|
));
|
||||||
@@ -630,13 +887,14 @@ impl MePool {
|
|||||||
payload.as_ref(),
|
payload.as_ref(),
|
||||||
proto_flags,
|
proto_flags,
|
||||||
tag.as_ref().map(|tag| tag.as_slice()),
|
tag.as_ref().map(|tag| tag.as_slice()),
|
||||||
|
Some(_permit),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(TrySendError::Closed(cmd)) => {
|
Err(TrySendError::Closed(cmd)) => {
|
||||||
let Some(payload) = proxy_req_payload_from_command(cmd) else {
|
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else {
|
||||||
return Err(ProxyError::Proxy(
|
return Err(ProxyError::Proxy(
|
||||||
"ME writer rejected unexpected command type".into(),
|
"ME writer rejected unexpected command type".into(),
|
||||||
));
|
));
|
||||||
@@ -653,6 +911,7 @@ impl MePool {
|
|||||||
payload.as_ref(),
|
payload.as_ref(),
|
||||||
proto_flags,
|
proto_flags,
|
||||||
tag.as_ref().map(|tag| tag.as_slice()),
|
tag.as_ref().map(|tag| tag.as_slice()),
|
||||||
|
Some(_permit),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
@@ -667,6 +926,7 @@ impl MePool {
|
|||||||
payload.as_ref(),
|
payload.as_ref(),
|
||||||
proto_flags,
|
proto_flags,
|
||||||
tag.as_ref().map(|tag| tag.as_slice()),
|
tag.as_ref().map(|tag| tag.as_slice()),
|
||||||
|
Some(_permit),
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
@@ -165,7 +170,7 @@ async fn recv_data_count(rx: &mut mpsc::Receiver<WriterCommand>, budget: Duratio
|
|||||||
while Instant::now().duration_since(start) < budget {
|
while Instant::now().duration_since(start) < budget {
|
||||||
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
|
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
|
||||||
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
|
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
|
||||||
Ok(Some(WriterCommand::Data(_))) => data_count += 1,
|
Ok(Some(WriterCommand::Data { .. })) => data_count += 1,
|
||||||
Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1,
|
Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1,
|
||||||
Ok(Some(WriterCommand::ProxyReq(_))) => data_count += 1,
|
Ok(Some(WriterCommand::ProxyReq(_))) => data_count += 1,
|
||||||
Ok(Some(WriterCommand::ControlAndFlush(_))) => data_count += 1,
|
Ok(Some(WriterCommand::ControlAndFlush(_))) => data_count += 1,
|
||||||
@@ -185,7 +190,7 @@ async fn recv_first_data_payload(
|
|||||||
while Instant::now().duration_since(start) < budget {
|
while Instant::now().duration_since(start) < budget {
|
||||||
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
|
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
|
||||||
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
|
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
|
||||||
Ok(Some(WriterCommand::Data(payload))) => return Some(payload.to_vec()),
|
Ok(Some(WriterCommand::Data { payload, .. })) => return Some(payload.to_vec()),
|
||||||
Ok(Some(WriterCommand::DataAndFlush(payload))) => return Some(payload.to_vec()),
|
Ok(Some(WriterCommand::DataAndFlush(payload))) => return Some(payload.to_vec()),
|
||||||
Ok(Some(_)) => {}
|
Ok(Some(_)) => {}
|
||||||
Ok(None) => break,
|
Ok(None) => break,
|
||||||
@@ -240,6 +245,7 @@ async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() {
|
|||||||
b"hello",
|
b"hello",
|
||||||
0,
|
0,
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -299,6 +305,7 @@ async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay
|
|||||||
b"storm",
|
b"storm",
|
||||||
0,
|
0,
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -356,6 +363,7 @@ async fn send_proxy_req_uses_writer_source_ip_when_advertised_our_addr_differs()
|
|||||||
b"route",
|
b"route",
|
||||||
0,
|
0,
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
@@ -411,6 +419,7 @@ async fn send_proxy_req_blocking_fallback_uses_writer_source_ip() {
|
|||||||
b"blocking",
|
b"blocking",
|
||||||
0,
|
0,
|
||||||
None,
|
None,
|
||||||
|
None,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
});
|
});
|
||||||
|
|||||||
Reference in New Issue
Block a user