mirror of https://github.com/telemt/telemt.git
287 lines
9.3 KiB
Rust
287 lines
9.3 KiB
Rust
use std::collections::HashMap;
|
|
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering};
|
|
use std::time::{Duration, Instant};
|
|
|
|
use tokio::sync::mpsc;
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
use super::codec::WriterCommand;
|
|
use super::pool::{MePool, MeWriter, WriterContour};
|
|
use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode};
|
|
use crate::crypto::SecureRandom;
|
|
use crate::network::probe::NetworkDecision;
|
|
use crate::stats::Stats;
|
|
|
|
async fn make_pool() -> (Arc<MePool>, Arc<SecureRandom>) {
|
|
let general = GeneralConfig {
|
|
me_route_no_writer_mode: MeRouteNoWriterMode::AsyncRecoveryFailfast,
|
|
me_route_no_writer_wait_ms: 50,
|
|
me_writer_pick_mode: MeWriterPickMode::SortedRr,
|
|
me_deterministic_writer_sort: true,
|
|
..GeneralConfig::default()
|
|
};
|
|
|
|
let rng = Arc::new(SecureRandom::new());
|
|
let pool = MePool::new(
|
|
None,
|
|
vec![1u8; 32],
|
|
None,
|
|
false,
|
|
None,
|
|
Vec::new(),
|
|
1,
|
|
None,
|
|
12,
|
|
1200,
|
|
HashMap::new(),
|
|
HashMap::new(),
|
|
None,
|
|
NetworkDecision::default(),
|
|
None,
|
|
rng.clone(),
|
|
Arc::new(Stats::default()),
|
|
general.me_keepalive_enabled,
|
|
general.me_keepalive_interval_secs,
|
|
general.me_keepalive_jitter_secs,
|
|
general.me_keepalive_payload_random,
|
|
general.rpc_proxy_req_every,
|
|
general.me_warmup_stagger_enabled,
|
|
general.me_warmup_step_delay_ms,
|
|
general.me_warmup_step_jitter_ms,
|
|
general.me_reconnect_max_concurrent_per_dc,
|
|
general.me_reconnect_backoff_base_ms,
|
|
general.me_reconnect_backoff_cap_ms,
|
|
general.me_reconnect_fast_retry_count,
|
|
general.me_single_endpoint_shadow_writers,
|
|
general.me_single_endpoint_outage_mode_enabled,
|
|
general.me_single_endpoint_outage_disable_quarantine,
|
|
general.me_single_endpoint_outage_backoff_min_ms,
|
|
general.me_single_endpoint_outage_backoff_max_ms,
|
|
general.me_single_endpoint_shadow_rotate_every_secs,
|
|
general.me_floor_mode,
|
|
general.me_adaptive_floor_idle_secs,
|
|
general.me_adaptive_floor_min_writers_single_endpoint,
|
|
general.me_adaptive_floor_min_writers_multi_endpoint,
|
|
general.me_adaptive_floor_recover_grace_secs,
|
|
general.me_adaptive_floor_writers_per_core_total,
|
|
general.me_adaptive_floor_cpu_cores_override,
|
|
general.me_adaptive_floor_max_extra_writers_single_per_core,
|
|
general.me_adaptive_floor_max_extra_writers_multi_per_core,
|
|
general.me_adaptive_floor_max_active_writers_per_core,
|
|
general.me_adaptive_floor_max_warm_writers_per_core,
|
|
general.me_adaptive_floor_max_active_writers_global,
|
|
general.me_adaptive_floor_max_warm_writers_global,
|
|
general.hardswap,
|
|
general.me_pool_drain_ttl_secs,
|
|
general.me_instadrain,
|
|
general.me_pool_drain_threshold,
|
|
general.me_pool_drain_soft_evict_enabled,
|
|
general.me_pool_drain_soft_evict_grace_secs,
|
|
general.me_pool_drain_soft_evict_per_writer,
|
|
general.me_pool_drain_soft_evict_budget_per_core,
|
|
general.me_pool_drain_soft_evict_cooldown_ms,
|
|
general.effective_me_pool_force_close_secs(),
|
|
general.me_pool_min_fresh_ratio,
|
|
general.me_hardswap_warmup_delay_min_ms,
|
|
general.me_hardswap_warmup_delay_max_ms,
|
|
general.me_hardswap_warmup_extra_passes,
|
|
general.me_hardswap_warmup_pass_backoff_base_ms,
|
|
general.me_bind_stale_mode,
|
|
general.me_bind_stale_ttl_secs,
|
|
general.me_secret_atomic_snapshot,
|
|
general.me_deterministic_writer_sort,
|
|
general.me_writer_pick_mode,
|
|
general.me_writer_pick_sample_size,
|
|
MeSocksKdfPolicy::default(),
|
|
general.me_writer_cmd_channel_capacity,
|
|
general.me_route_channel_capacity,
|
|
general.me_route_backpressure_base_timeout_ms,
|
|
general.me_route_backpressure_high_timeout_ms,
|
|
general.me_route_backpressure_high_watermark_pct,
|
|
general.me_reader_route_data_wait_ms,
|
|
general.me_health_interval_ms_unhealthy,
|
|
general.me_health_interval_ms_healthy,
|
|
general.me_warn_rate_limit_ms,
|
|
general.me_route_no_writer_mode,
|
|
general.me_route_no_writer_wait_ms,
|
|
general.me_route_hybrid_max_wait_ms,
|
|
general.me_route_blocking_send_timeout_ms,
|
|
general.me_route_inline_recovery_attempts,
|
|
general.me_route_inline_recovery_wait_ms,
|
|
);
|
|
|
|
(pool, rng)
|
|
}
|
|
|
|
async fn insert_writer(
|
|
pool: &Arc<MePool>,
|
|
writer_id: u64,
|
|
writer_dc: i32,
|
|
addr: SocketAddr,
|
|
register_in_registry: bool,
|
|
) -> mpsc::Receiver<WriterCommand> {
|
|
let (tx, rx) = mpsc::channel::<WriterCommand>(8);
|
|
let writer = MeWriter {
|
|
id: writer_id,
|
|
addr,
|
|
source_ip: addr.ip(),
|
|
writer_dc,
|
|
generation: pool.current_generation(),
|
|
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
|
|
created_at: Instant::now(),
|
|
tx: tx.clone(),
|
|
cancel: CancellationToken::new(),
|
|
degraded: Arc::new(AtomicBool::new(false)),
|
|
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
|
|
draining: Arc::new(AtomicBool::new(false)),
|
|
draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)),
|
|
drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)),
|
|
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
|
|
};
|
|
|
|
pool.writers.write().await.push(writer);
|
|
{
|
|
let mut map = pool.proxy_map_v4.write().await;
|
|
map.entry(writer_dc)
|
|
.or_insert_with(Vec::new)
|
|
.push((addr.ip(), addr.port()));
|
|
}
|
|
pool.rebuild_endpoint_dc_map().await;
|
|
if register_in_registry {
|
|
pool.registry.register_writer(writer_id, tx).await;
|
|
}
|
|
rx
|
|
}
|
|
|
|
async fn recv_data_count(rx: &mut mpsc::Receiver<WriterCommand>, budget: Duration) -> usize {
|
|
let start = Instant::now();
|
|
let mut data_count = 0usize;
|
|
while Instant::now().duration_since(start) < budget {
|
|
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
|
|
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
|
|
Ok(Some(WriterCommand::Data(_))) => data_count += 1,
|
|
Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1,
|
|
Ok(Some(WriterCommand::Close)) => {}
|
|
Ok(None) => break,
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
data_count
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() {
|
|
let (pool, _rng) = make_pool().await;
|
|
pool.rr.store(0, Ordering::Relaxed);
|
|
|
|
let (conn_id, _rx) = pool.registry.register().await;
|
|
let mut stale_rx = insert_writer(
|
|
&pool,
|
|
10,
|
|
2,
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 10)), 443),
|
|
false,
|
|
)
|
|
.await;
|
|
let mut live_rx = insert_writer(
|
|
&pool,
|
|
11,
|
|
2,
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 11)), 443),
|
|
true,
|
|
)
|
|
.await;
|
|
|
|
let result = pool
|
|
.send_proxy_req(
|
|
conn_id,
|
|
2,
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 30000),
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443),
|
|
b"hello",
|
|
0,
|
|
None,
|
|
)
|
|
.await;
|
|
|
|
assert!(result.is_ok());
|
|
assert_eq!(
|
|
recv_data_count(&mut stale_rx, Duration::from_millis(50)).await,
|
|
0
|
|
);
|
|
assert_eq!(
|
|
recv_data_count(&mut live_rx, Duration::from_millis(50)).await,
|
|
1
|
|
);
|
|
|
|
let bound = pool.registry.get_writer(conn_id).await;
|
|
assert!(bound.is_some());
|
|
assert_eq!(bound.expect("writer should be bound").writer_id, 11);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay() {
|
|
let (pool, _rng) = make_pool().await;
|
|
pool.rr.store(0, Ordering::Relaxed);
|
|
|
|
let (conn_id, _rx) = pool.registry.register().await;
|
|
|
|
let mut stale_rx_1 = insert_writer(
|
|
&pool,
|
|
21,
|
|
2,
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 21)), 443),
|
|
false,
|
|
)
|
|
.await;
|
|
let mut stale_rx_2 = insert_writer(
|
|
&pool,
|
|
22,
|
|
2,
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 22)), 443),
|
|
false,
|
|
)
|
|
.await;
|
|
let mut live_rx = insert_writer(
|
|
&pool,
|
|
23,
|
|
2,
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 1, 23)), 443),
|
|
true,
|
|
)
|
|
.await;
|
|
|
|
let result = pool
|
|
.send_proxy_req(
|
|
conn_id,
|
|
2,
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 30001),
|
|
SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443),
|
|
b"storm",
|
|
0,
|
|
None,
|
|
)
|
|
.await;
|
|
|
|
assert!(result.is_ok());
|
|
assert_eq!(
|
|
recv_data_count(&mut stale_rx_1, Duration::from_millis(50)).await,
|
|
0
|
|
);
|
|
assert_eq!(
|
|
recv_data_count(&mut stale_rx_2, Duration::from_millis(50)).await,
|
|
0
|
|
);
|
|
assert_eq!(
|
|
recv_data_count(&mut live_rx, Duration::from_millis(50)).await,
|
|
1
|
|
);
|
|
|
|
let writers = pool.writers.read().await;
|
|
let writer_ids = writers.iter().map(|w| w.id).collect::<Vec<_>>();
|
|
drop(writers);
|
|
assert_eq!(writer_ids, vec![23]);
|
|
}
|