feat(proxy,failover,transport): Hyper 上游客户端精细计时、连续失败退避与连接泄漏修复

Proxy:
- 将上游 HTTP 客户端从 reqwest 替换为 hyper,新增 InstrumentedConnector
  实现 TCP 连接/TLS 握手级别的独立计时,上报 connection_reused 等指标
- 前端展示细粒度代理计时(连接复用、等待响应头等)

Failover:
- 引入连续失败退避机制,每 10 次失败递增退避间隔
- 检测 H2 max outbound streams 错误并触发上游客户端重建
- 新增 HTTPClientPool.reset_upstream_client 支持按需重建缓存客户端

连接泄漏修复:
- Handler 异常路径确保 response_ctx 被正确关闭
- HubResponseStream 迭代结束后在 finally 块中清理 stream_id
- HubTunnelTransport.handle_request 捕获所有异常并清理流状态
This commit is contained in:
fawney19
2026-03-06 17:55:14 +08:00
parent 2269617a9f
commit 050cba9563
15 changed files with 902 additions and 72 deletions

View File

@@ -10,7 +10,7 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
[[package]] [[package]]
name = "aether-proxy" name = "aether-proxy"
version = "0.2.2" version = "0.2.3"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"arc-swap", "arc-swap",
@@ -21,6 +21,9 @@ dependencies = [
"flate2", "flate2",
"futures-util", "futures-util",
"hex", "hex",
"http-body-util",
"hyper",
"hyper-util",
"libc", "libc",
"ratatui", "ratatui",
"reqwest", "reqwest",
@@ -33,8 +36,10 @@ dependencies = [
"tar", "tar",
"thiserror 2.0.18", "thiserror 2.0.18",
"tokio", "tokio",
"tokio-rustls",
"tokio-tungstenite", "tokio-tungstenite",
"toml", "toml",
"tower-service",
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
"url", "url",

View File

@@ -7,7 +7,11 @@ description = "Tunnel proxy for Aether"
[dependencies] [dependencies]
tokio = { version = "1", features = ["full"] } tokio = { version = "1", features = ["full"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream", "http2"] } reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream", "http2"] }
hyper = { version = "1", features = ["client", "http1", "http2"] }
hyper-util = { version = "0.1", features = ["client", "client-legacy", "http1", "http2", "tokio"] }
http-body-util = "0.1"
tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] } tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] }
tokio-rustls = "0.26"
futures-util = "0.3" futures-util = "0.3"
base64 = "0.22" base64 = "0.22"
clap = { version = "4", features = ["derive", "env"] } clap = { version = "4", features = ["derive", "env"] }
@@ -31,6 +35,7 @@ libc = "0.2"
flate2 = "1" flate2 = "1"
tar = "0.4" tar = "0.4"
socket2 = { version = "0.5", features = ["all"] } socket2 = { version = "0.5", features = ["all"] }
tower-service = "0.3"
webpki-roots = "0.26" webpki-roots = "0.26"
[profile.release] [profile.release]

View File

@@ -13,8 +13,8 @@ use crate::config::{Config, ServerEntry};
use crate::net; use crate::net;
use crate::registration::client::AetherClient; use crate::registration::client::AetherClient;
use crate::runtime::{self, DynamicConfig}; use crate::runtime::{self, DynamicConfig};
use crate::safe_dns::SafeDnsResolver;
use crate::state::{AppState, ProxyMetrics, ServerContext}; use crate::state::{AppState, ProxyMetrics, ServerContext};
use crate::upstream_client;
use crate::{hardware, target_filter, tunnel}; use crate::{hardware, target_filter, tunnel};
/// Run the full application lifecycle after config has been parsed. /// Run the full application lifecycle after config has been parsed.
@@ -67,27 +67,10 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
config.dns_cache_capacity, config.dns_cache_capacity,
)); ));
// Build reqwest client for tunnel upstream requests (shared). // Build Hyper client for tunnel upstream requests (shared).
// Inject SafeDnsResolver so reqwest only connects to addresses that were // DNS still flows through validated addresses from DnsCache, while the
// validated by validate_target() — this eliminates the DNS rebinding // custom connector exposes per-request connect/TLS timing when available.
// TOCTTOU gap where a second DNS lookup could return a private IP. let upstream_client = upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
let safe_resolver = SafeDnsResolver::new(Arc::clone(&dns_cache));
let mut reqwest_builder = reqwest::Client::builder()
.dns_resolver(Arc::new(safe_resolver))
.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host)
.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs))
.connect_timeout(Duration::from_secs(config.upstream_connect_timeout_secs))
.tcp_nodelay(config.upstream_tcp_nodelay);
if config.upstream_tcp_keepalive_secs > 0 {
reqwest_builder = reqwest_builder.tcp_keepalive(Some(Duration::from_secs(
config.upstream_tcp_keepalive_secs,
)));
}
let reqwest_client = reqwest_builder
.build()
.expect("failed to build reqwest client");
// Register with each Aether server and build per-server contexts. // Register with each Aether server and build per-server contexts.
// Wrapped in Arc<Mutex> so retry_failed_registrations can append later. // Wrapped in Arc<Mutex> so retry_failed_registrations can append later.
@@ -160,7 +143,7 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
let state = Arc::new(AppState { let state = Arc::new(AppState {
config: Arc::new(config), config: Arc::new(config),
dns_cache, dns_cache,
reqwest_client, upstream_client,
tunnel_tls_config, tunnel_tls_config,
}); });

View File

@@ -4,11 +4,11 @@ mod hardware;
mod net; mod net;
mod registration; mod registration;
mod runtime; mod runtime;
mod safe_dns;
mod setup; mod setup;
mod state; mod state;
mod target_filter; mod target_filter;
mod tunnel; mod tunnel;
mod upstream_client;
use std::path::PathBuf; use std::path::PathBuf;

View File

@@ -8,14 +8,15 @@ use crate::config::Config;
use crate::registration::client::AetherClient; use crate::registration::client::AetherClient;
use crate::runtime::SharedDynamicConfig; use crate::runtime::SharedDynamicConfig;
use crate::target_filter::DnsCache; use crate::target_filter::DnsCache;
use crate::upstream_client::UpstreamClient;
/// Central application state shared across all servers/tunnels. /// Central application state shared across all servers/tunnels.
pub struct AppState { pub struct AppState {
pub config: Arc<Config>, pub config: Arc<Config>,
/// DNS cache for upstream target resolution (shared). /// DNS cache for upstream target resolution (shared).
pub dns_cache: Arc<DnsCache>, pub dns_cache: Arc<DnsCache>,
/// Reqwest client for tunnel upstream requests (shared). /// Hyper client for tunnel upstream requests with validated DNS and connection timing.
pub reqwest_client: reqwest::Client, pub upstream_client: UpstreamClient,
/// Shared TLS config for tunnel WebSocket connections (avoids re-parsing root CAs on each reconnect). /// Shared TLS config for tunnel WebSocket connections (avoids re-parsing root CAs on each reconnect).
pub tunnel_tls_config: Arc<rustls::ClientConfig>, pub tunnel_tls_config: Arc<rustls::ClientConfig>,
} }

View File

@@ -9,11 +9,13 @@ use std::time::{Duration, Instant};
use bytes::Bytes; use bytes::Bytes;
use futures_util::StreamExt; use futures_util::StreamExt;
use http_body_util::BodyExt;
use tokio::sync::mpsc; use tokio::sync::mpsc;
use tracing::{debug, warn}; use tracing::{debug, warn};
use crate::state::{AppState, ServerContext}; use crate::state::{AppState, ServerContext};
use crate::target_filter; use crate::target_filter;
use crate::upstream_client::{self, UpstreamRequestBody};
use super::protocol::{ use super::protocol::{
compress_payload, decompress_if_gzip, flags, Frame, MsgType, RequestMeta, ResponseMeta, compress_payload, decompress_if_gzip, flags, Frame, MsgType, RequestMeta, ResponseMeta,
@@ -203,44 +205,61 @@ async fn handle_stream_inner(
let dns_ms = connect_start.elapsed().as_millis() as u64; let dns_ms = connect_start.elapsed().as_millis() as u64;
// Execute upstream request // Execute upstream request
let client = &state.reqwest_client; let client = &state.upstream_client;
let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS)); let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
let method: reqwest::Method = meta.method.parse().unwrap_or(reqwest::Method::GET); let method: hyper::Method = meta.method.parse().unwrap_or(hyper::Method::GET);
// Build a complete HeaderMap from tunnel headers, then set it all at once let mut request = match hyper::Request::builder()
// via .headers() which *replaces* reqwest defaults (e.g. Accept: */*), .method(method)
// ensuring upstream sees exactly what Aether server intended. .uri(meta.url.as_str())
let mut header_map = reqwest::header::HeaderMap::with_capacity(meta.headers.len()); .body(UpstreamRequestBody::new(body.clone()))
{
Ok(request) => request,
Err(e) => {
send_error(
frame_tx,
stream_id,
&format!("invalid upstream request: {e}"),
)
.await;
return None;
}
};
let headers = request.headers_mut();
for (k, v) in &meta.headers { for (k, v) in &meta.headers {
let k_lower = k.to_ascii_lowercase(); let k_lower = k.to_ascii_lowercase();
if BLOCKED_HEADERS.contains(&k_lower.as_str()) { if BLOCKED_HEADERS.contains(&k_lower.as_str()) {
continue; continue;
} }
if let (Ok(name), Ok(value)) = ( if let (Ok(name), Ok(value)) = (
reqwest::header::HeaderName::from_bytes(k.as_bytes()), hyper::header::HeaderName::from_bytes(k.as_bytes()),
reqwest::header::HeaderValue::from_str(v), hyper::header::HeaderValue::from_str(v),
) { ) {
header_map.insert(name, value); headers.insert(name, value);
} }
} }
let mut req = client.request(method, &meta.url).headers(header_map);
let body_size = body.len(); let body_size = body.len();
if !body.is_empty() { let mut captured_connection = upstream_client::capture_connection(&mut request);
req = req.body(body); let connection_start = Instant::now();
} let connection_capture = tokio::spawn(async move {
req = req.timeout(timeout); let connected = captured_connection.wait_for_connection_metadata().await;
connected
.as_ref()
.map(|_| connection_start.elapsed().as_millis() as u64)
});
let upstream_start = Instant::now(); let upstream_start = Instant::now();
let response = match req.send().await { let response = match tokio::time::timeout(timeout, client.request(request)).await {
Ok(r) => r, Ok(Ok(response)) => response,
Err(e) => { Ok(Err(e)) => {
connection_capture.abort();
server server
.metrics .metrics
.failed_requests .failed_requests
.fetch_add(1, Ordering::Release); .fetch_add(1, Ordering::Release);
let msg = if e.is_timeout() { let msg = if e.is_connect() {
"upstream timeout".to_string()
} else if e.is_connect() {
format!("upstream connect error: {e}") format!("upstream connect error: {e}")
} else { } else {
format!("upstream error: {e}") format!("upstream error: {e}")
@@ -248,6 +267,15 @@ async fn handle_stream_inner(
send_error(frame_tx, stream_id, &msg).await; send_error(frame_tx, stream_id, &msg).await;
return None; return None;
} }
Err(_) => {
connection_capture.abort();
server
.metrics
.failed_requests
.fetch_add(1, Ordering::Release);
send_error(frame_tx, stream_id, "upstream timeout").await;
return None;
}
}; };
// Capture connection-establishment duration (DNS + TCP/TLS + TTFB) // Capture connection-establishment duration (DNS + TCP/TLS + TTFB)
@@ -257,18 +285,34 @@ async fn handle_stream_inner(
// Send RESPONSE_HEADERS // Send RESPONSE_HEADERS
let status = response.status().as_u16(); let status = response.status().as_u16();
let ttfb_ms = upstream_start.elapsed().as_millis() as u64; let ttfb_ms = upstream_start.elapsed().as_millis() as u64;
// Short timeout: on connection reuse hyper may never fire the connect
// callback, so avoid blocking indefinitely.
let connection_acquire_ms =
match tokio::time::timeout(Duration::from_millis(100), connection_capture).await {
Ok(Ok(ms)) => ms,
Ok(Err(_)) => None, // JoinError (task panicked / cancelled)
Err(_) => None, // timeout -- task is detached but lightweight
};
let request_timing =
upstream_client::resolve_request_timing(&response, connection_acquire_ms, ttfb_ms);
let mut resp_headers: Vec<(String, String)> = Vec::with_capacity(response.headers().len() + 1); let mut resp_headers: Vec<(String, String)> = Vec::with_capacity(response.headers().len() + 1);
for (k, v) in response.headers() { for (k, v) in response.headers() {
if let Ok(vs) = v.to_str() { if let Ok(vs) = v.to_str() {
resp_headers.push((k.as_str().to_string(), vs.to_string())); resp_headers.push((k.as_str().to_string(), vs.to_string()));
} }
} }
// Inject proxy timing (same format as delegate mode)
let timing = serde_json::json!({ let timing = serde_json::json!({
"dns_ms": dns_ms, "dns_ms": dns_ms,
"connection_acquire_ms": request_timing.connection_acquire_ms,
"connection_reused": request_timing.connection_reused,
"connect_ms": request_timing.connect_ms,
"tls_ms": request_timing.tls_ms,
"ttfb_ms": ttfb_ms, "ttfb_ms": ttfb_ms,
"upstream_ms": ttfb_ms, "upstream_ms": ttfb_ms,
"upstream_processing_ms": ttfb_ms.saturating_sub(dns_ms), "response_wait_ms": request_timing.response_wait_ms,
"upstream_processing_ms": request_timing.response_wait_ms,
"timing_source": "instrumented_connector",
"total_ms": connect_elapsed.as_millis() as u64,
"body_size": body_size, "body_size": body_size,
"mode": "tunnel", "mode": "tunnel",
}); });
@@ -298,7 +342,7 @@ async fn handle_stream_inner(
// (e.g. uncompressed SSE text). Already-compressed data (gzip/br from // (e.g. uncompressed SSE text). Already-compressed data (gzip/br from
// upstream Content-Encoding) won't shrink further and will be sent as-is // upstream Content-Encoding) won't shrink further and will be sent as-is
// thanks to the size check in compress_payload(). // thanks to the size check in compress_payload().
let mut stream = response.bytes_stream(); let mut stream = response.into_body().into_data_stream();
while let Some(chunk_result) = stream.next().await { while let Some(chunk_result) = stream.next().await {
match chunk_result { match chunk_result {
Ok(chunk) => { Ok(chunk) => {

View File

@@ -0,0 +1,434 @@
use std::future::Future;
use std::io;
use std::net::IpAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use bytes::Bytes;
use http_body_util::Full;
use hyper::rt;
use hyper::Response;
use hyper::Uri;
pub use hyper_util::client::legacy::connect::capture_connection;
use hyper_util::client::legacy::connect::dns::Name;
use hyper_util::client::legacy::connect::{Connected, Connection, HttpConnector};
use hyper_util::client::legacy::Client;
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
use rustls::pki_types::ServerName;
use rustls::ClientConfig;
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;
use tower_service::Service;
use crate::config::Config;
use crate::target_filter::{self, DnsCache};
type BoxError = Box<dyn std::error::Error + Send + Sync>;
type PlainStream = TokioIo<TcpStream>;
type TlsStream = TokioIo<tokio_rustls::client::TlsStream<TcpStream>>;
pub type UpstreamRequestBody = Full<Bytes>;
pub type UpstreamClient = Client<InstrumentedConnector, UpstreamRequestBody>;
#[derive(Clone, Copy, Debug, Default)]
pub struct ConnectTiming {
pub connect_ms: u64,
pub tls_ms: u64,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct RequestTiming {
pub connection_acquire_ms: u64,
pub connect_ms: u64,
pub tls_ms: u64,
pub response_wait_ms: u64,
pub connection_reused: bool,
}
#[derive(Clone)]
pub struct ValidatedResolver {
dns_cache: Arc<DnsCache>,
}
impl ValidatedResolver {
pub fn new(dns_cache: Arc<DnsCache>) -> Self {
Self { dns_cache }
}
}
pub struct ValidatedAddrs {
inner: std::vec::IntoIter<std::net::SocketAddr>,
}
impl Iterator for ValidatedAddrs {
type Item = std::net::SocketAddr;
fn next(&mut self) -> Option<Self::Item> {
self.inner.next()
}
}
impl Service<Name> for ValidatedResolver {
type Response = ValidatedAddrs;
type Error = io::Error;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, name: Name) -> Self::Future {
let dns_cache = Arc::clone(&self.dns_cache);
let host = name.as_str().to_string();
Box::pin(async move {
if let Some(addrs) = dns_cache.get_by_host(&host).await {
return Ok(ValidatedAddrs {
inner: (*addrs).clone().into_iter(),
});
}
let resolved = target_filter::resolve_public_addrs(&host, 0, dns_cache.as_ref())
.await
.map_err(|err| io::Error::other(err.to_string()))?;
Ok(ValidatedAddrs {
inner: resolved.into_iter(),
})
})
}
}
#[derive(Clone)]
pub struct InstrumentedConnector {
http: HttpConnector<ValidatedResolver>,
tls_config: Arc<ClientConfig>,
}
impl Service<Uri> for InstrumentedConnector {
type Response = TimedConn;
type Error = BoxError;
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.http.poll_ready(cx).map_err(Into::into)
}
fn call(&mut self, dst: Uri) -> Self::Future {
let scheme = dst.scheme_str().map(|value| value.to_ascii_lowercase());
let tls_config = Arc::clone(&self.tls_config);
let connecting = self.http.call(dst.clone());
let connect_start = std::time::Instant::now();
Box::pin(async move {
match scheme.as_deref() {
Some("http") => {
let tcp = connecting.await.map_err(|err| Box::new(err) as BoxError)?;
let connect_ms = connect_start.elapsed().as_millis() as u64;
Ok(TimedConn::new(
MaybeHttpsStream::Http(tcp),
ConnectTiming {
connect_ms,
tls_ms: 0,
},
))
}
Some("https") => {
let server_name = resolve_server_name(&dst)?;
let tcp = connecting.await.map_err(|err| Box::new(err) as BoxError)?;
let connect_ms = connect_start.elapsed().as_millis() as u64;
let tls_start = std::time::Instant::now();
let tls_stream = TlsConnector::from(tls_config)
.connect(server_name, tcp.into_inner())
.await
.map_err(io::Error::other)?;
let tls_ms = tls_start.elapsed().as_millis() as u64;
Ok(TimedConn::new(
MaybeHttpsStream::Https(TokioIo::new(tls_stream)),
ConnectTiming { connect_ms, tls_ms },
))
}
Some(other) => Err(io::Error::other(format!("unsupported scheme {other}")).into()),
None => Err(io::Error::other("missing scheme").into()),
}
})
}
}
pub fn build_upstream_client(config: &Config, dns_cache: Arc<DnsCache>) -> UpstreamClient {
let mut http = HttpConnector::new_with_resolver(ValidatedResolver::new(dns_cache));
http.enforce_http(false);
http.set_connect_timeout(Some(Duration::from_secs(
config.upstream_connect_timeout_secs,
)));
http.set_nodelay(config.upstream_tcp_nodelay);
if config.upstream_tcp_keepalive_secs > 0 {
http.set_keepalive(Some(Duration::from_secs(
config.upstream_tcp_keepalive_secs,
)));
} else {
http.set_keepalive(None);
}
let connector = InstrumentedConnector {
http,
tls_config: build_tls_config(),
};
let mut builder = Client::builder(TokioExecutor::new());
builder.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host);
builder.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs));
builder.pool_timer(TokioTimer::new());
builder.build(connector)
}
pub fn resolve_request_timing<B>(
response: &Response<B>,
connection_acquire_ms: Option<u64>,
ttfb_ms: u64,
) -> RequestTiming {
let raw = response
.extensions()
.get::<ConnectTiming>()
.copied()
.unwrap_or_default();
let raw_connection_ms = raw.connect_ms.saturating_add(raw.tls_ms);
let measured_acquire_ms = connection_acquire_ms.unwrap_or(raw_connection_ms.min(ttfb_ms));
let likely_reused = measured_acquire_ms <= 5 && raw_connection_ms > 0;
let connector_matches_request = raw_connection_ms <= measured_acquire_ms.saturating_add(25);
let (connect_ms, tls_ms) = if likely_reused || !connector_matches_request {
(0, 0)
} else {
(raw.connect_ms, raw.tls_ms)
};
RequestTiming {
connection_acquire_ms: measured_acquire_ms,
connect_ms,
tls_ms,
response_wait_ms: ttfb_ms.saturating_sub(measured_acquire_ms),
connection_reused: likely_reused,
}
}
fn build_tls_config() -> Arc<ClientConfig> {
let root_store =
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
let mut config = ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
Arc::new(config)
}
fn resolve_server_name(uri: &Uri) -> Result<ServerName<'static>, BoxError> {
let host = uri.host().ok_or_else(|| io::Error::other("missing host"))?;
let host = host.trim_start_matches('[').trim_end_matches(']');
if let Ok(ip) = host.parse::<IpAddr>() {
return Ok(ServerName::from(ip));
}
Ok(ServerName::try_from(host.to_string())?)
}
pub struct TimedConn {
inner: MaybeHttpsStream,
timing: ConnectTiming,
}
impl TimedConn {
fn new(inner: MaybeHttpsStream, timing: ConnectTiming) -> Self {
Self { inner, timing }
}
}
impl Connection for TimedConn {
fn connected(&self) -> Connected {
self.inner.connected().extra(self.timing)
}
}
impl rt::Read for TimedConn {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: rt::ReadBufCursor<'_>,
) -> Poll<Result<(), io::Error>> {
Pin::new(&mut self.inner).poll_read(cx, buf)
}
}
impl rt::Write for TimedConn {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, io::Error>> {
Pin::new(&mut self.inner).poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
Pin::new(&mut self.inner).poll_flush(cx)
}
fn poll_shutdown(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), io::Error>> {
Pin::new(&mut self.inner).poll_shutdown(cx)
}
fn is_write_vectored(&self) -> bool {
self.inner.is_write_vectored()
}
fn poll_write_vectored(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> Poll<Result<usize, io::Error>> {
Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
}
}
pub enum MaybeHttpsStream {
Http(PlainStream),
Https(TlsStream),
}
impl Connection for MaybeHttpsStream {
fn connected(&self) -> Connected {
match self {
Self::Http(stream) => stream.connected(),
Self::Https(stream) => {
let (tcp, tls) = stream.inner().get_ref();
if tls.alpn_protocol() == Some(b"h2") {
tcp.connected().negotiated_h2()
} else {
tcp.connected()
}
}
}
}
}
impl rt::Read for MaybeHttpsStream {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: rt::ReadBufCursor<'_>,
) -> Poll<Result<(), io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_read(cx, buf),
Self::Https(stream) => Pin::new(stream).poll_read(cx, buf),
}
}
}
impl rt::Write for MaybeHttpsStream {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_write(cx, buf),
Self::Https(stream) => Pin::new(stream).poll_write(cx, buf),
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_flush(cx),
Self::Https(stream) => Pin::new(stream).poll_flush(cx),
}
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_shutdown(cx),
Self::Https(stream) => Pin::new(stream).poll_shutdown(cx),
}
}
fn is_write_vectored(&self) -> bool {
match self {
Self::Http(stream) => stream.is_write_vectored(),
Self::Https(stream) => stream.is_write_vectored(),
}
}
fn poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> Poll<Result<usize, io::Error>> {
match Pin::get_mut(self) {
Self::Http(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
Self::Https(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use hyper::Response;
#[test]
fn fresh_connection_uses_connector_breakdown() {
let mut response = Response::new(());
response.extensions_mut().insert(ConnectTiming {
connect_ms: 80,
tls_ms: 40,
});
let timing = resolve_request_timing(&response, Some(125), 600);
assert_eq!(timing.connection_acquire_ms, 125);
assert_eq!(timing.connect_ms, 80);
assert_eq!(timing.tls_ms, 40);
assert_eq!(timing.response_wait_ms, 475);
assert!(!timing.connection_reused);
}
#[test]
fn reused_connection_zeroes_stale_connect_timings() {
let mut response = Response::new(());
response.extensions_mut().insert(ConnectTiming {
connect_ms: 70,
tls_ms: 30,
});
let timing = resolve_request_timing(&response, Some(0), 310);
assert_eq!(timing.connection_acquire_ms, 0);
assert_eq!(timing.connect_ms, 0);
assert_eq!(timing.tls_ms, 0);
assert_eq!(timing.response_wait_ms, 310);
assert!(timing.connection_reused);
}
#[test]
fn falls_back_to_connector_timings_when_capture_missing() {
let mut response = Response::new(());
response.extensions_mut().insert(ConnectTiming {
connect_ms: 55,
tls_ms: 25,
});
let timing = resolve_request_timing(&response, None, 400);
assert_eq!(timing.connection_acquire_ms, 80);
assert_eq!(timing.connect_ms, 55);
assert_eq!(timing.tls_ms, 25);
assert_eq!(timing.response_wait_ms, 320);
assert!(!timing.connection_reused);
}
}

View File

@@ -566,7 +566,12 @@ const proxyTimingBreakdown = (proxy: Record<string, unknown>): string => {
} }
const ttfbMs = t.ttfb_ms ?? t.upstream_ms const ttfbMs = t.ttfb_ms ?? t.upstream_ms
const processingMs = t.upstream_processing_ms ?? ( const responseWaitMs = t.response_wait_ms ?? (
t.connection_acquire_ms != null && ttfbMs != null
? Math.max(0, (ttfbMs as number) - (t.connection_acquire_ms as number))
: null
)
const legacyWaitMs = t.upstream_processing_ms ?? (
ttfbMs != null && t.connect_ms != null && t.tls_ms != null ttfbMs != null && t.connect_ms != null && t.tls_ms != null
? Math.max(0, (ttfbMs as number) - (t.connect_ms as number) - (t.tls_ms as number)) ? Math.max(0, (ttfbMs as number) - (t.connect_ms as number) - (t.tls_ms as number))
: null : null
@@ -575,6 +580,9 @@ const proxyTimingBreakdown = (proxy: Record<string, unknown>): string => {
if (t.dns_ms != null && (t.dns_ms as number) > 0) { if (t.dns_ms != null && (t.dns_ms as number) > 0) {
parts.push(`DNS ${formatLatency(t.dns_ms as number)}`) parts.push(`DNS ${formatLatency(t.dns_ms as number)}`)
} }
if (t.connection_reused === true) {
parts.push('复用连接')
}
if (t.connect_ms != null && (t.connect_ms as number) > 0) { if (t.connect_ms != null && (t.connect_ms as number) > 0) {
parts.push(`连接 ${formatLatency(t.connect_ms as number)}`) parts.push(`连接 ${formatLatency(t.connect_ms as number)}`)
} }
@@ -584,8 +592,10 @@ const proxyTimingBreakdown = (proxy: Record<string, unknown>): string => {
if (ttfbMs != null && (ttfbMs as number) > 0) { if (ttfbMs != null && (ttfbMs as number) > 0) {
parts.push(`TTFB ${formatLatency(ttfbMs as number)}`) parts.push(`TTFB ${formatLatency(ttfbMs as number)}`)
} }
if (processingMs != null && (processingMs as number) > 0) { if (responseWaitMs != null && (responseWaitMs as number) > 0) {
parts.push(`上游处理 ${formatLatency(Math.round(processingMs as number))}`) parts.push(`等待响应头 ${formatLatency(Math.round(responseWaitMs as number))}`)
} else if (legacyWaitMs != null && (legacyWaitMs as number) > 0) {
parts.push(`等待响应头(旧版估算) ${formatLatency(Math.round(legacyWaitMs as number))}`)
} }
// 计算 Aether→代理 之间无法解释的耗时差 // 计算 Aether→代理 之间无法解释的耗时差

View File

@@ -1287,6 +1287,14 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
response_ctx = None response_ctx = None
continue continue
try:
if response_ctx is not None:
await response_ctx.__aexit__(None, None, None)
except Exception:
pass
finally:
response_ctx = None
from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor
error_text = await ChatSyncExecutor(self)._extract_error_text(e) error_text = await ChatSyncExecutor(self)._extract_error_text(e)
@@ -1306,6 +1314,13 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
raise raise
except Exception: except Exception:
try:
if response_ctx is not None:
await response_ctx.__aexit__(None, None, None)
except Exception:
pass
finally:
response_ctx = None
raise raise
# 类型断言:成功执行后这些变量不会为 None # 类型断言:成功执行后这些变量不会为 None

View File

@@ -800,6 +800,14 @@ class CliStreamMixin:
response_ctx = None response_ctx = None
continue continue
try:
if response_ctx is not None:
await response_ctx.__aexit__(None, None, None)
except Exception:
pass
finally:
response_ctx = None
error_text = await self._extract_error_text(e) error_text = await self._extract_error_text(e)
logger.error( logger.error(
f"Provider 返回错误状态: {e.response.status_code}\n Response: {error_text}" f"Provider 返回错误状态: {e.response.status_code}\n Response: {error_text}"
@@ -818,6 +826,13 @@ class CliStreamMixin:
raise raise
except Exception: except Exception:
try:
if response_ctx is not None:
await response_ctx.__aexit__(None, None, None)
except Exception:
pass
finally:
response_ctx = None
raise raise
# 类型断言:成功执行后这些变量不会为 None # 类型断言:成功执行后这些变量不会为 None

View File

@@ -458,6 +458,67 @@ class HTTPClientPool:
return httpx.AsyncClient(**client_config) # type: ignore[arg-type] return httpx.AsyncClient(**client_config) # type: ignore[arg-type]
@classmethod @classmethod
async def reset_upstream_client(
cls,
delegate_cfg: dict[str, Any] | None,
proxy_config: dict[str, Any] | None = None,
tls_profile: str | None = None,
) -> bool:
"""Reset cached upstream client for the given proxy/tunnel route.
Returns True when a cached client was closed and removed.
For the shared no-proxy default client this is a no-op to avoid
disrupting unrelated in-flight requests.
"""
if delegate_cfg and delegate_cfg.get("tunnel"):
node_id = str(delegate_cfg.get("node_id") or "")
if not node_id:
return False
lock = cls._get_proxy_clients_lock()
async with lock:
client = cls._tunnel_clients.pop(node_id, None)
if client is None:
return False
try:
await client.aclose()
except Exception as exc:
logger.warning("关闭 Tunnel 客户端失败(node_id={}): {}", node_id, exc)
return True
if not proxy_config:
proxy_config = get_system_proxy_config()
base_cache_key = compute_proxy_cache_key(proxy_config)
if base_cache_key == "__no_proxy__":
return False
cache_key_prefixes = [base_cache_key]
tls_profile_key = str(tls_profile or "").strip().lower()
if tls_profile_key:
cache_key_prefixes = [f"{base_cache_key}::tls:{tls_profile_key}"]
lock = cls._get_proxy_clients_lock()
async with lock:
keys_to_remove = [
key
for key in list(cls._proxy_clients.keys())
if any(
key == prefix
or key.startswith(f"{prefix}::")
or key.startswith(f"{prefix}::tls:")
for prefix in cache_key_prefixes
)
]
clients = [cls._proxy_clients.pop(key)[0] for key in keys_to_remove]
for client in clients:
try:
await client.aclose()
except Exception as exc:
logger.warning("关闭上游代理客户端失败: {}", exc)
return bool(clients)
async def get_upstream_client( async def get_upstream_client(
cls, cls,
delegate_cfg: dict[str, Any] | None, delegate_cfg: dict[str, Any] | None,

View File

@@ -49,6 +49,11 @@ class FailoverEngine:
# Hard constraint: streaming first chunk probe timeout # Hard constraint: streaming first chunk probe timeout
STREAM_FIRST_CHUNK_TIMEOUT_SECONDS: int = 30 STREAM_FIRST_CHUNK_TIMEOUT_SECONDS: int = 30
RETRY_BACKOFF_EVERY_FAILURES: int = 10
RETRY_ROTATE_CLIENT_EVERY_FAILURES: int = 40
RETRY_BACKOFF_BASE_SECONDS: float = 0.025
RETRY_BACKOFF_MAX_SECONDS: float = 0.15
STREAM_CAPACITY_BACKOFF_SECONDS: float = 0.3
def __init__( def __init__(
self, self,
@@ -61,6 +66,121 @@ class FailoverEngine:
self._error_classifier = error_classifier or ErrorClassifier(db=db) self._error_classifier = error_classifier or ErrorClassifier(db=db)
self._recorder = recorder or CandidateRecorder(db) self._recorder = recorder or CandidateRecorder(db)
@staticmethod
def _collect_error_messages(error: Exception | None) -> str:
if error is None:
return ""
parts: list[str] = []
for item in (
getattr(error, "message", None),
getattr(error, "upstream_response", None),
str(error),
):
if isinstance(item, str) and item.strip():
parts.append(item.strip())
cause = getattr(error, "cause", None)
if cause is not None and cause is not error:
for item in (
getattr(cause, "message", None),
getattr(cause, "upstream_response", None),
str(cause),
):
if isinstance(item, str) and item.strip():
parts.append(item.strip())
return " | ".join(parts)
@classmethod
def _is_stream_capacity_error(cls, error: Exception | None) -> bool:
lowered = cls._collect_error_messages(error).lower()
return (
"max outbound streams" in lowered
or "too many concurrent streams" in lowered
or "max concurrent streams" in lowered
)
@classmethod
def _compute_retry_backoff_seconds(
cls,
*,
consecutive_failures: int,
error: Exception | None,
) -> float:
if consecutive_failures <= 0:
return 0.0
if cls._is_stream_capacity_error(error):
return cls.STREAM_CAPACITY_BACKOFF_SECONDS
if consecutive_failures % cls.RETRY_BACKOFF_EVERY_FAILURES != 0:
return 0.0
step = max(1, consecutive_failures // cls.RETRY_BACKOFF_EVERY_FAILURES)
return min(cls.RETRY_BACKOFF_BASE_SECONDS * step, cls.RETRY_BACKOFF_MAX_SECONDS)
@classmethod
def _should_rotate_upstream_client(
cls,
*,
consecutive_failures: int,
error: Exception | None,
) -> bool:
if cls._is_stream_capacity_error(error):
return True
return (
consecutive_failures >= cls.RETRY_ROTATE_CLIENT_EVERY_FAILURES
and consecutive_failures % cls.RETRY_ROTATE_CLIENT_EVERY_FAILURES == 0
)
async def _rotate_upstream_client(self, candidate: ProviderCandidate) -> bool:
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
resolve_delegate_config,
resolve_effective_proxy,
)
effective_proxy = resolve_effective_proxy(
getattr(candidate.provider, "proxy", None),
getattr(candidate.key, "proxy", None),
)
delegate_cfg = resolve_delegate_config(effective_proxy)
return await HTTPClientPool.reset_upstream_client(
delegate_cfg, proxy_config=effective_proxy
)
async def _apply_retry_pacing(
self,
*,
candidate: ProviderCandidate,
consecutive_failures: int,
error: Exception | None,
request_id: str | None,
) -> None:
should_rotate = self._should_rotate_upstream_client(
consecutive_failures=consecutive_failures,
error=error,
)
if should_rotate:
rotated = await self._rotate_upstream_client(candidate)
if rotated:
logger.warning(
" [{}] 连续失败 {} 次,已重建上游客户端复用",
request_id,
consecutive_failures,
)
backoff_seconds = self._compute_retry_backoff_seconds(
consecutive_failures=consecutive_failures,
error=error,
)
if backoff_seconds > 0:
logger.warning(
" [{}] 连续失败 {} 次,退避 {:.0f}ms 后继续尝试",
request_id,
consecutive_failures,
backoff_seconds * 1000,
)
await asyncio.sleep(backoff_seconds)
async def execute( async def execute(
self, self,
*, *,
@@ -95,6 +215,7 @@ class FailoverEngine:
candidates = candidates[:max_candidates] candidates = candidates[:max_candidates]
attempt_count = 0 attempt_count = 0
consecutive_failures = 0
last_status_code: int | None = None last_status_code: int | None = None
# For logging / dispatcher parity only; callers may pass an exact value. # For logging / dispatcher parity only; callers may pass an exact value.
@@ -142,20 +263,23 @@ class FailoverEngine:
continue continue
if isinstance(candidate, PoolCandidate): if isinstance(candidate, PoolCandidate):
pool_result, attempt_count, last_status_code = await self._execute_pool_candidate( pool_result, attempt_count, consecutive_failures, last_status_code = (
candidate=candidate, await self._execute_pool_candidate(
candidate_index=candidate_index, candidate=candidate,
attempt_func=attempt_func, candidate_index=candidate_index,
retry_policy=retry_policy, attempt_func=attempt_func,
request_id=request_id, retry_policy=retry_policy,
user_id=user_id, request_id=request_id,
api_key_id=api_key_id, user_id=user_id,
candidate_record_map=candidate_record_map, api_key_id=api_key_id,
candidate_keys_fallback=candidate_keys_fallback, candidate_record_map=candidate_record_map,
candidates=candidates, candidate_keys_fallback=candidate_keys_fallback,
attempt_count=attempt_count, candidates=candidates,
max_attempts=max_attempts, attempt_count=attempt_count,
execution_error_handler=execution_error_handler, max_attempts=max_attempts,
execution_error_handler=execution_error_handler,
consecutive_failures=consecutive_failures,
)
) )
if pool_result is not None: if pool_result is not None:
return pool_result return pool_result
@@ -217,6 +341,7 @@ class FailoverEngine:
retry_policy=retry_policy, retry_policy=retry_policy,
) )
consecutive_failures = 0
return ExecutionResult( return ExecutionResult(
success=True, success=True,
attempt_result=attempt_result, attempt_result=attempt_result,
@@ -240,6 +365,13 @@ class FailoverEngine:
last_status_code = exc.http_status last_status_code = exc.http_status
self._record_attempt_failure(record_id, exc, exc.http_status) self._record_attempt_failure(record_id, exc, exc.http_status)
action = FailoverAction.CONTINUE action = FailoverAction.CONTINUE
consecutive_failures += 1
await self._apply_retry_pacing(
candidate=candidate,
consecutive_failures=consecutive_failures,
error=exc,
request_id=request_id,
)
except Exception as exc: except Exception as exc:
outcome = await self._handle_attempt_error( outcome = await self._handle_attempt_error(
@@ -263,6 +395,14 @@ class FailoverEngine:
max_retries = outcome.max_retries max_retries = outcome.max_retries
if outcome.stop_result is not None: if outcome.stop_result is not None:
return outcome.stop_result return outcome.stop_result
if action in {FailoverAction.CONTINUE, FailoverAction.RETRY}:
consecutive_failures += 1
await self._apply_retry_pacing(
candidate=candidate,
consecutive_failures=consecutive_failures,
error=exc,
request_id=request_id,
)
# action switch: continue/ retry # action switch: continue/ retry
if action == FailoverAction.CONTINUE: if action == FailoverAction.CONTINUE:
@@ -313,9 +453,10 @@ class FailoverEngine:
candidate_keys_fallback: list[CandidateKey], candidate_keys_fallback: list[CandidateKey],
candidates: list[ProviderCandidate], candidates: list[ProviderCandidate],
attempt_count: int, attempt_count: int,
consecutive_failures: int,
max_attempts: int | None, max_attempts: int | None,
execution_error_handler: Any, execution_error_handler: Any,
) -> tuple[ExecutionResult | None, int, int | None]: ) -> tuple[ExecutionResult | None, int, int, int | None]:
"""Execute a PoolCandidate with in-pool key failover.""" """Execute a PoolCandidate with in-pool key failover."""
last_status_code: int | None = None last_status_code: int | None = None
retry_slots_per_key = self._get_pool_key_max_retries(candidate, retry_policy) retry_slots_per_key = self._get_pool_key_max_retries(candidate, retry_policy)
@@ -421,6 +562,7 @@ class FailoverEngine:
retry_policy=retry_policy, retry_policy=retry_policy,
) )
consecutive_failures = 0
return ( return (
ExecutionResult( ExecutionResult(
success=True, success=True,
@@ -441,6 +583,7 @@ class FailoverEngine:
request_candidate_id=record_id, request_candidate_id=record_id,
), ),
attempt_count, attempt_count,
consecutive_failures,
last_status_code, last_status_code,
) )
@@ -448,6 +591,13 @@ class FailoverEngine:
last_status_code = exc.http_status last_status_code = exc.http_status
self._record_attempt_failure(record_id, exc, exc.http_status) self._record_attempt_failure(record_id, exc, exc.http_status)
action = FailoverAction.CONTINUE action = FailoverAction.CONTINUE
consecutive_failures += 1
await self._apply_retry_pacing(
candidate=candidate,
consecutive_failures=consecutive_failures,
error=exc,
request_id=request_id,
)
except Exception as exc: except Exception as exc:
outcome = await self._handle_pool_attempt_error( outcome = await self._handle_pool_attempt_error(
@@ -465,6 +615,14 @@ class FailoverEngine:
action = outcome.action action = outcome.action
last_status_code = outcome.last_status_code last_status_code = outcome.last_status_code
max_retries_for_key = min(outcome.max_retries, retry_slots_per_key) max_retries_for_key = min(outcome.max_retries, retry_slots_per_key)
if action in {FailoverAction.CONTINUE, FailoverAction.RETRY}:
consecutive_failures += 1
await self._apply_retry_pacing(
candidate=candidate,
consecutive_failures=consecutive_failures,
error=exc,
request_id=request_id,
)
if action == FailoverAction.CONTINUE: if action == FailoverAction.CONTINUE:
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map: if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
@@ -496,9 +654,9 @@ class FailoverEngine:
from_retry_idx=composite_retry_index + 1, from_retry_idx=composite_retry_index + 1,
retry_policy=retry_policy, retry_policy=retry_policy,
) )
return None, attempt_count, last_status_code return None, attempt_count, consecutive_failures, last_status_code
return None, attempt_count, last_status_code return None, attempt_count, consecutive_failures, last_status_code
async def _handle_pool_attempt_error( async def _handle_pool_attempt_error(
self, self,

View File

@@ -642,6 +642,9 @@ class HubTunnelTransport(httpx.AsyncBaseTransport):
) )
self._cleanup_stream(manager, stream_state) self._cleanup_stream(manager, stream_state)
raise httpx.ReadTimeout("hub tunnel request timeout") from None raise httpx.ReadTimeout("hub tunnel request timeout") from None
except Exception:
self._cleanup_stream(manager, stream_state)
raise
def _cleanup_stream( def _cleanup_stream(
self, self,
@@ -665,8 +668,11 @@ class HubResponseStream(httpx.AsyncByteStream):
self._timeout = timeout self._timeout = timeout
async def __aiter__(self) -> AsyncGenerator[bytes, None]: async def __aiter__(self) -> AsyncGenerator[bytes, None]:
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout): try:
yield chunk async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
yield chunk
finally:
self._manager.remove_stream(self._stream_state.stream_id)
async def aclose(self) -> None: async def aclose(self) -> None:
self._manager.remove_stream(self._stream_state.stream_id) self._manager.remove_stream(self._stream_state.stream_id)

View File

@@ -456,3 +456,49 @@ async def test_error_stop_pattern_without_status_codes_matches_any() -> None:
# No status_codes filter, pattern matches -> stop # No status_codes filter, pattern matches -> stop
assert result.success is False assert result.success is False
assert attempt.await_count == 1 assert attempt.await_count == 1
@pytest.mark.asyncio
async def test_failover_engine_applies_backoff_every_tenth_failure(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
engine = FailoverEngine(db, error_classifier=_StubErrorClassifier(action=ErrorAction.BREAK))
sleep_mock = AsyncMock()
rotate_mock = AsyncMock(return_value=False)
monkeypatch.setattr("src.services.candidate.failover.asyncio.sleep", sleep_mock)
monkeypatch.setattr(engine, "_rotate_upstream_client", rotate_mock)
await engine._apply_retry_pacing(
candidate=_make_candidate(),
consecutive_failures=10,
error=RuntimeError("boom"),
request_id="req-1",
)
sleep_mock.assert_awaited_once()
rotate_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_failover_engine_rotates_client_on_stream_capacity_error(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
engine = FailoverEngine(db, error_classifier=_StubErrorClassifier(action=ErrorAction.BREAK))
sleep_mock = AsyncMock()
rotate_mock = AsyncMock(return_value=True)
monkeypatch.setattr("src.services.candidate.failover.asyncio.sleep", sleep_mock)
monkeypatch.setattr(engine, "_rotate_upstream_client", rotate_mock)
await engine._apply_retry_pacing(
candidate=_make_candidate(),
consecutive_failures=3,
error=RuntimeError("LocalProtocolError: Max outbound streams is 100, 100 open"),
request_id="req-2",
)
rotate_mock.assert_awaited_once()
sleep_mock.assert_awaited_once()

View File

@@ -0,0 +1,47 @@
from __future__ import annotations
import pytest
from src.services.proxy_node.hub_transport import HubResponseStream
from src.services.proxy_node.tunnel_manager import _StreamState
class _Manager:
def __init__(self) -> None:
self.removed: list[int] = []
def remove_stream(self, stream_id: int) -> None:
self.removed.append(stream_id)
@pytest.mark.asyncio
async def test_hub_response_stream_removes_stream_after_normal_iteration() -> None:
manager = _Manager()
state = _StreamState(7)
state.set_response_headers(200, {})
state.push_body_chunk(b"hello")
state.set_done()
stream = HubResponseStream(manager, state, timeout=0.1)
chunks = []
async for chunk in stream:
chunks.append(chunk)
assert chunks == [b"hello"]
assert manager.removed == [7]
@pytest.mark.asyncio
async def test_hub_response_stream_removes_stream_after_body_error() -> None:
manager = _Manager()
state = _StreamState(9)
state.set_response_headers(200, {})
state.set_error("boom")
stream = HubResponseStream(manager, state, timeout=0.1)
with pytest.raises(Exception):
async for _chunk in stream:
pass
assert manager.removed == [9]