[network] in upstream

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-02-18 06:02:24 +03:00
parent eb9ac7fae4
commit 67bae1cf2a

View File

@@ -355,6 +355,8 @@ impl UpstreamManager {
&self, &self,
prefer_ipv6: bool, prefer_ipv6: bool,
dc_overrides: &HashMap<String, Vec<String>>, dc_overrides: &HashMap<String, Vec<String>>,
ipv4_enabled: bool,
ipv6_enabled: bool,
) -> Vec<StartupPingResult> { ) -> Vec<StartupPingResult> {
let upstreams: Vec<(usize, UpstreamConfig)> = { let upstreams: Vec<(usize, UpstreamConfig)> = {
let guard = self.upstreams.read().await; let guard = self.upstreams.read().await;
@@ -374,10 +376,8 @@ impl UpstreamManager {
UpstreamType::Socks5 { address, .. } => format!("socks5://{}", address), UpstreamType::Socks5 { address, .. } => format!("socks5://{}", address),
}; };
let mut v6_results = Vec::new(); let mut v6_results = Vec::with_capacity(NUM_DCS);
let mut v4_results = Vec::new(); if ipv6_enabled {
// === Ping IPv6 first ===
for dc_zero_idx in 0..NUM_DCS { for dc_zero_idx in 0..NUM_DCS {
let dc_v6 = TG_DATACENTERS_V6[dc_zero_idx]; let dc_v6 = TG_DATACENTERS_V6[dc_zero_idx];
let addr_v6 = SocketAddr::new(dc_v6, TG_DATACENTER_PORT); let addr_v6 = SocketAddr::new(dc_v6, TG_DATACENTER_PORT);
@@ -415,8 +415,20 @@ impl UpstreamManager {
}; };
v6_results.push(ping_result); v6_results.push(ping_result);
} }
} else {
for dc_zero_idx in 0..NUM_DCS {
let dc_v6 = TG_DATACENTERS_V6[dc_zero_idx];
v6_results.push(DcPingResult {
dc_idx: dc_zero_idx + 1,
dc_addr: SocketAddr::new(dc_v6, TG_DATACENTER_PORT),
rtt_ms: None,
error: Some("ipv6 disabled".to_string()),
});
}
}
// === Then ping IPv4 === let mut v4_results = Vec::with_capacity(NUM_DCS);
if ipv4_enabled {
for dc_zero_idx in 0..NUM_DCS { for dc_zero_idx in 0..NUM_DCS {
let dc_v4 = TG_DATACENTERS_V4[dc_zero_idx]; let dc_v4 = TG_DATACENTERS_V4[dc_zero_idx];
let addr_v4 = SocketAddr::new(dc_v4, TG_DATACENTER_PORT); let addr_v4 = SocketAddr::new(dc_v4, TG_DATACENTER_PORT);
@@ -454,6 +466,17 @@ impl UpstreamManager {
}; };
v4_results.push(ping_result); v4_results.push(ping_result);
} }
} else {
for dc_zero_idx in 0..NUM_DCS {
let dc_v4 = TG_DATACENTERS_V4[dc_zero_idx];
v4_results.push(DcPingResult {
dc_idx: dc_zero_idx + 1,
dc_addr: SocketAddr::new(dc_v4, TG_DATACENTER_PORT),
rtt_ms: None,
error: Some("ipv4 disabled".to_string()),
});
}
}
// === Ping DC overrides (v4/v6) === // === Ping DC overrides (v4/v6) ===
for (dc_key, addrs) in dc_overrides { for (dc_key, addrs) in dc_overrides {
@@ -470,6 +493,9 @@ impl UpstreamManager {
match addr_str.parse::<SocketAddr>() { match addr_str.parse::<SocketAddr>() {
Ok(addr) => { Ok(addr) => {
let is_v6 = addr.is_ipv6(); let is_v6 = addr.is_ipv6();
if (is_v6 && !ipv6_enabled) || (!is_v6 && !ipv4_enabled) {
continue;
}
let result = tokio::time::timeout( let result = tokio::time::timeout(
Duration::from_secs(DC_PING_TIMEOUT_SECS), Duration::from_secs(DC_PING_TIMEOUT_SECS),
self.ping_single_dc(&upstream_config, addr) self.ping_single_dc(&upstream_config, addr)
@@ -551,7 +577,7 @@ impl UpstreamManager {
/// Background health check: rotates through DCs, 30s interval. /// Background health check: rotates through DCs, 30s interval.
/// Uses preferred IP version based on config. /// Uses preferred IP version based on config.
pub async fn run_health_checks(&self, prefer_ipv6: bool) { pub async fn run_health_checks(&self, prefer_ipv6: bool, ipv4_enabled: bool, ipv6_enabled: bool) {
let mut dc_rotation = 0usize; let mut dc_rotation = 0usize;
loop { loop {
@@ -560,16 +586,24 @@ impl UpstreamManager {
let dc_zero_idx = dc_rotation % NUM_DCS; let dc_zero_idx = dc_rotation % NUM_DCS;
dc_rotation += 1; dc_rotation += 1;
let dc_addr = if prefer_ipv6 { let primary_v6 = SocketAddr::new(TG_DATACENTERS_V6[dc_zero_idx], TG_DATACENTER_PORT);
SocketAddr::new(TG_DATACENTERS_V6[dc_zero_idx], TG_DATACENTER_PORT) let primary_v4 = SocketAddr::new(TG_DATACENTERS_V4[dc_zero_idx], TG_DATACENTER_PORT);
let dc_addr = if prefer_ipv6 && ipv6_enabled {
primary_v6
} else if ipv4_enabled {
primary_v4
} else if ipv6_enabled {
primary_v6
} else { } else {
SocketAddr::new(TG_DATACENTERS_V4[dc_zero_idx], TG_DATACENTER_PORT) continue;
}; };
let fallback_addr = if prefer_ipv6 { let fallback_addr = if dc_addr.is_ipv6() && ipv4_enabled {
SocketAddr::new(TG_DATACENTERS_V4[dc_zero_idx], TG_DATACENTER_PORT) Some(primary_v4)
} else if dc_addr.is_ipv4() && ipv6_enabled {
Some(primary_v6)
} else { } else {
SocketAddr::new(TG_DATACENTERS_V6[dc_zero_idx], TG_DATACENTER_PORT) None
}; };
let count = self.upstreams.read().await.len(); let count = self.upstreams.read().await.len();
@@ -608,6 +642,7 @@ impl UpstreamManager {
// Try fallback // Try fallback
debug!(dc = dc_zero_idx + 1, "Health check failed, trying fallback"); debug!(dc = dc_zero_idx + 1, "Health check failed, trying fallback");
if let Some(fallback_addr) = fallback_addr {
let start2 = Instant::now(); let start2 = Instant::now();
let result2 = tokio::time::timeout( let result2 = tokio::time::timeout(
Duration::from_secs(10), Duration::from_secs(10),
@@ -652,6 +687,17 @@ impl UpstreamManager {
} }
} }
u.last_check = std::time::Instant::now(); u.last_check = std::time::Instant::now();
continue;
}
let mut guard = self.upstreams.write().await;
let u = &mut guard[i];
u.fails += 1;
if u.fails > 3 {
u.healthy = false;
warn!("Upstream unhealthy (no fallback family)");
}
u.last_check = std::time::Instant::now();
} }
} }
} }