mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
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:
7
aether-proxy/Cargo.lock
generated
7
aether-proxy/Cargo.lock
generated
@@ -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",
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
@@ -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,
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|
||||||
|
|||||||
@@ -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>,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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) => {
|
||||||
|
|||||||
434
aether-proxy/src/upstream_client.rs
Normal file
434
aether-proxy/src/upstream_client.rs
Normal 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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→代理 之间无法解释的耗时差
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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,7 +263,8 @@ 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 = (
|
||||||
|
await self._execute_pool_candidate(
|
||||||
candidate=candidate,
|
candidate=candidate,
|
||||||
candidate_index=candidate_index,
|
candidate_index=candidate_index,
|
||||||
attempt_func=attempt_func,
|
attempt_func=attempt_func,
|
||||||
@@ -156,6 +278,8 @@ class FailoverEngine:
|
|||||||
attempt_count=attempt_count,
|
attempt_count=attempt_count,
|
||||||
max_attempts=max_attempts,
|
max_attempts=max_attempts,
|
||||||
execution_error_handler=execution_error_handler,
|
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,
|
||||||
|
|||||||
@@ -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]:
|
||||||
|
try:
|
||||||
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
|
async for chunk in self._stream_state.iter_body(chunk_timeout=self._timeout):
|
||||||
yield chunk
|
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)
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
47
tests/unit/test_hub_transport_stream_cleanup.py
Normal file
47
tests/unit/test_hub_transport_stream_cleanup.py
Normal 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]
|
||||||
Reference in New Issue
Block a user