mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix(tunnel): gateway 回复 WebSocket Pong 防止 proxy stale 断连
gateway reader 之前用 `_ => {}` 忽略了 proxy 发来的 WebSocket Ping,
导致 proxy 的 stale_timeout (5s) 持续触发断连重连循环,
所有隧道请求都因等待 response headers 超时而失败。
同时补充 hub 诊断日志和连续请求集成测试。
This commit is contained in:
@@ -241,6 +241,7 @@ mod tests {
|
|||||||
use super::build_direct_execution_frame_stream;
|
use super::build_direct_execution_frame_stream;
|
||||||
use crate::execution_runtime::transport::{
|
use crate::execution_runtime::transport::{
|
||||||
execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime,
|
execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime,
|
||||||
|
DirectUpstreamResponse,
|
||||||
};
|
};
|
||||||
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
|
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
@@ -584,4 +585,172 @@ mod tests {
|
|||||||
"local tunnel path should preserve the original proxy error text"
|
"local tunnel path should preserve the original proxy error text"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn second_local_tunnel_request_works_after_first_completes() {
|
||||||
|
let state = AppState::new().expect("app state should build");
|
||||||
|
let tunnel_app = state.tunnel.app_state();
|
||||||
|
let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8);
|
||||||
|
let (proxy_close_tx, _) = watch::channel(false);
|
||||||
|
tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new(
|
||||||
|
900,
|
||||||
|
"node-1".to_string(),
|
||||||
|
"Node 1".to_string(),
|
||||||
|
proxy_tx,
|
||||||
|
proxy_close_tx,
|
||||||
|
16,
|
||||||
|
)));
|
||||||
|
|
||||||
|
let plan = ExecutionPlan {
|
||||||
|
request_id: "req-reuse-1".into(),
|
||||||
|
candidate_id: Some("cand-reuse-1".into()),
|
||||||
|
provider_name: Some("openai".into()),
|
||||||
|
provider_id: "prov-1".into(),
|
||||||
|
endpoint_id: "ep-1".into(),
|
||||||
|
key_id: "key-1".into(),
|
||||||
|
method: "POST".into(),
|
||||||
|
url: "https://example.com/chat".into(),
|
||||||
|
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
|
||||||
|
content_type: Some("application/json".into()),
|
||||||
|
content_encoding: None,
|
||||||
|
body: RequestBody::from_json(serde_json::json!({"stream": true})),
|
||||||
|
stream: true,
|
||||||
|
client_api_format: "openai:chat".into(),
|
||||||
|
provider_api_format: "openai:chat".into(),
|
||||||
|
model_name: Some("gpt-5".into()),
|
||||||
|
proxy: Some(tunnel_proxy_snapshot("http://127.0.0.1:1".to_string())),
|
||||||
|
tls_profile: None,
|
||||||
|
timeouts: Some(ExecutionTimeouts {
|
||||||
|
connect_ms: Some(5_000),
|
||||||
|
total_ms: Some(5_000),
|
||||||
|
..ExecutionTimeouts::default()
|
||||||
|
}),
|
||||||
|
};
|
||||||
|
|
||||||
|
// --- First request ---
|
||||||
|
let state1 = state.clone();
|
||||||
|
let plan1 = plan.clone();
|
||||||
|
let exec1 = tokio::spawn(async move {
|
||||||
|
execute_stream_plan_via_local_tunnel(&state1, &plan1).await
|
||||||
|
});
|
||||||
|
|
||||||
|
// Read request frames from proxy side
|
||||||
|
let req1_headers = match proxy_rx.recv().await.expect("req1 headers") {
|
||||||
|
Message::Binary(data) => data,
|
||||||
|
other => panic!("unexpected: {other:?}"),
|
||||||
|
};
|
||||||
|
let req1_header = tunnel_protocol::FrameHeader::parse(&req1_headers)
|
||||||
|
.expect("req1 header parse");
|
||||||
|
let _req1_body = proxy_rx.recv().await.expect("req1 body");
|
||||||
|
|
||||||
|
// Simulate proxy response
|
||||||
|
let resp_meta = serde_json::to_vec(&tunnel_protocol::ResponseMeta {
|
||||||
|
status: 200,
|
||||||
|
headers: vec![("content-type".to_string(), "text/event-stream".to_string())],
|
||||||
|
})
|
||||||
|
.unwrap();
|
||||||
|
let mut resp_headers = tunnel_protocol::encode_frame(
|
||||||
|
req1_header.stream_id,
|
||||||
|
tunnel_protocol::RESPONSE_HEADERS,
|
||||||
|
0,
|
||||||
|
&resp_meta,
|
||||||
|
);
|
||||||
|
tunnel_app
|
||||||
|
.hub
|
||||||
|
.handle_proxy_frame(900, &mut resp_headers)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let execution1 = exec1
|
||||||
|
.await
|
||||||
|
.expect("task")
|
||||||
|
.expect("transport")
|
||||||
|
.expect("execution");
|
||||||
|
|
||||||
|
// Consume the body stream fully
|
||||||
|
let mut resp1 = match execution1.response {
|
||||||
|
DirectUpstreamResponse::LocalTunnel(r) => r,
|
||||||
|
_ => panic!("expected local tunnel response"),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Send body + STREAM_END
|
||||||
|
let mut body_frame = tunnel_protocol::encode_frame(
|
||||||
|
req1_header.stream_id,
|
||||||
|
tunnel_protocol::RESPONSE_BODY,
|
||||||
|
0,
|
||||||
|
b"data: hello\n\n",
|
||||||
|
);
|
||||||
|
tunnel_app
|
||||||
|
.hub
|
||||||
|
.handle_proxy_frame(900, &mut body_frame)
|
||||||
|
.await;
|
||||||
|
let mut end_frame = tunnel_protocol::encode_frame(
|
||||||
|
req1_header.stream_id,
|
||||||
|
tunnel_protocol::STREAM_END,
|
||||||
|
0,
|
||||||
|
&[],
|
||||||
|
);
|
||||||
|
tunnel_app
|
||||||
|
.hub
|
||||||
|
.handle_proxy_frame(900, &mut end_frame)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
// Drain the body
|
||||||
|
while let Ok(Some(_)) = resp1.next_chunk().await {}
|
||||||
|
drop(resp1);
|
||||||
|
|
||||||
|
// --- Second request ---
|
||||||
|
let state2 = state.clone();
|
||||||
|
let plan2 = ExecutionPlan {
|
||||||
|
request_id: "req-reuse-2".into(),
|
||||||
|
candidate_id: Some("cand-reuse-2".into()),
|
||||||
|
..plan.clone()
|
||||||
|
};
|
||||||
|
let exec2 = tokio::spawn(async move {
|
||||||
|
execute_stream_plan_via_local_tunnel(&state2, &plan2).await
|
||||||
|
});
|
||||||
|
|
||||||
|
// Read second request's frames
|
||||||
|
let req2_headers = tokio::time::timeout(Duration::from_secs(2), proxy_rx.recv())
|
||||||
|
.await
|
||||||
|
.expect("second request should arrive within 2s")
|
||||||
|
.expect("req2 headers");
|
||||||
|
let req2_data = match req2_headers {
|
||||||
|
Message::Binary(data) => data,
|
||||||
|
other => panic!("unexpected: {other:?}"),
|
||||||
|
};
|
||||||
|
let req2_header = tunnel_protocol::FrameHeader::parse(&req2_data)
|
||||||
|
.expect("req2 header parse");
|
||||||
|
assert_eq!(req2_header.msg_type, tunnel_protocol::REQUEST_HEADERS);
|
||||||
|
|
||||||
|
// Simulate proxy response for second request
|
||||||
|
let mut resp2_headers = tunnel_protocol::encode_frame(
|
||||||
|
req2_header.stream_id,
|
||||||
|
tunnel_protocol::RESPONSE_HEADERS,
|
||||||
|
0,
|
||||||
|
&resp_meta,
|
||||||
|
);
|
||||||
|
tunnel_app
|
||||||
|
.hub
|
||||||
|
.handle_proxy_frame(900, &mut resp2_headers)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let execution2 = exec2
|
||||||
|
.await
|
||||||
|
.expect("task")
|
||||||
|
.expect("transport")
|
||||||
|
.expect("second execution should succeed");
|
||||||
|
assert_eq!(execution2.status_code, 200);
|
||||||
|
|
||||||
|
// Clean up
|
||||||
|
let mut end2 = tunnel_protocol::encode_frame(
|
||||||
|
req2_header.stream_id,
|
||||||
|
tunnel_protocol::STREAM_END,
|
||||||
|
0,
|
||||||
|
&[],
|
||||||
|
);
|
||||||
|
tunnel_app
|
||||||
|
.hub
|
||||||
|
.handle_proxy_frame(900, &mut end2)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -472,11 +472,21 @@ impl HubRouter {
|
|||||||
fn get_proxy_conn(&self, node_id: &str) -> Option<Arc<ProxyConn>> {
|
fn get_proxy_conn(&self, node_id: &str) -> Option<Arc<ProxyConn>> {
|
||||||
let map = self.proxy_conns.read();
|
let map = self.proxy_conns.read();
|
||||||
let conns = map.get(node_id)?;
|
let conns = map.get(node_id)?;
|
||||||
conns
|
let result = conns
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|c| c.is_available())
|
.filter(|c| c.is_available())
|
||||||
.min_by_key(|c| c.stream_count.load(Ordering::Relaxed))
|
.min_by_key(|c| c.stream_count.load(Ordering::Relaxed))
|
||||||
.cloned()
|
.cloned();
|
||||||
|
if result.is_none() && !conns.is_empty() {
|
||||||
|
warn!(
|
||||||
|
node_id = %node_id,
|
||||||
|
total_conns = conns.len(),
|
||||||
|
closing = conns.iter().filter(|c| c.outbound.is_closing()).count(),
|
||||||
|
draining = conns.iter().filter(|c| c.is_draining()).count(),
|
||||||
|
"no available proxy connection despite registered connections"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
result
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn has_local_proxy(&self, node_id: &str) -> bool {
|
pub fn has_local_proxy(&self, node_id: &str) -> bool {
|
||||||
@@ -531,7 +541,17 @@ impl HubRouter {
|
|||||||
self.proxy_to_local
|
self.proxy_to_local
|
||||||
.insert((proxy_conn.id, proxy_stream_id), local_stream_id);
|
.insert((proxy_conn.id, proxy_stream_id), local_stream_id);
|
||||||
|
|
||||||
match proxy_conn.send(Message::Binary(header_frame.into())) {
|
let send_status = proxy_conn.send(Message::Binary(header_frame.into()));
|
||||||
|
debug!(
|
||||||
|
node_id = %node_id,
|
||||||
|
conn_id = proxy_conn.id,
|
||||||
|
proxy_stream_id = proxy_stream_id,
|
||||||
|
local_stream_id = local_stream_id,
|
||||||
|
stream_count = proxy_conn.stream_count.load(Ordering::Relaxed),
|
||||||
|
send_status = ?send_status,
|
||||||
|
"open_local_stream dispatched"
|
||||||
|
);
|
||||||
|
match send_status {
|
||||||
SendStatus::Queued => Ok(local_stream),
|
SendStatus::Queued => Ok(local_stream),
|
||||||
SendStatus::Closed | SendStatus::Congested => {
|
SendStatus::Closed | SendStatus::Congested => {
|
||||||
self.cleanup_local_stream(local_stream_id);
|
self.cleanup_local_stream(local_stream_id);
|
||||||
@@ -1082,4 +1102,88 @@ mod tests {
|
|||||||
|
|
||||||
assert!(proxy_rx.try_recv().is_err());
|
assert!(proxy_rx.try_recv().is_err());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn second_stream_works_after_first_completes_via_stream_end() {
|
||||||
|
let hub = HubRouter::new(ControlPlaneClient::disabled());
|
||||||
|
|
||||||
|
let (proxy_tx, mut proxy_rx) = bounded_queue(8);
|
||||||
|
let (proxy_close_tx, _) = watch::channel(false);
|
||||||
|
let proxy = Arc::new(ProxyConn::new(
|
||||||
|
400,
|
||||||
|
"node-reuse".to_string(),
|
||||||
|
"Node Reuse".to_string(),
|
||||||
|
proxy_tx,
|
||||||
|
proxy_close_tx,
|
||||||
|
16,
|
||||||
|
));
|
||||||
|
hub.register_proxy(Arc::clone(&proxy));
|
||||||
|
|
||||||
|
// First request: open stream, send body, simulate proxy response + STREAM_END
|
||||||
|
let stream1 = hub
|
||||||
|
.open_local_stream("node-reuse", &build_meta())
|
||||||
|
.expect("open first stream");
|
||||||
|
let _ = proxy_rx.try_recv().expect("first headers frame");
|
||||||
|
hub.push_local_request_body(stream1.id, Bytes::new(), true)
|
||||||
|
.expect("first body");
|
||||||
|
let _ = proxy_rx.try_recv().expect("first body frame");
|
||||||
|
|
||||||
|
// Simulate proxy sending RESPONSE_HEADERS
|
||||||
|
let resp_meta = serde_json::to_vec(&serde_json::json!({
|
||||||
|
"status": 200,
|
||||||
|
"headers": []
|
||||||
|
}))
|
||||||
|
.unwrap();
|
||||||
|
let mut resp_headers_frame = protocol::encode_frame(
|
||||||
|
// Extract the proxy_stream_id from the request headers frame
|
||||||
|
2, // first stream_id allocated
|
||||||
|
protocol::RESPONSE_HEADERS,
|
||||||
|
0,
|
||||||
|
&resp_meta,
|
||||||
|
);
|
||||||
|
hub.handle_proxy_frame(400, &mut resp_headers_frame).await;
|
||||||
|
|
||||||
|
// Simulate proxy sending STREAM_END
|
||||||
|
let mut end_frame = protocol::encode_frame(2, protocol::STREAM_END, 0, &[]);
|
||||||
|
hub.handle_proxy_frame(400, &mut end_frame).await;
|
||||||
|
|
||||||
|
// Verify stream_count went back to 0
|
||||||
|
assert_eq!(
|
||||||
|
proxy.stream_count.load(std::sync::atomic::Ordering::Relaxed),
|
||||||
|
0,
|
||||||
|
"stream_count should be 0 after STREAM_END"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Second request: should work
|
||||||
|
let stream2 = hub
|
||||||
|
.open_local_stream("node-reuse", &build_meta())
|
||||||
|
.expect("open second stream should succeed");
|
||||||
|
let second_headers = proxy_rx.try_recv().expect("second headers frame");
|
||||||
|
let second_data = match second_headers {
|
||||||
|
Message::Binary(data) => data.to_vec(),
|
||||||
|
other => panic!("unexpected message: {other:?}"),
|
||||||
|
};
|
||||||
|
let header =
|
||||||
|
protocol::FrameHeader::parse(&second_data).expect("second request header frame");
|
||||||
|
assert_eq!(header.msg_type, protocol::REQUEST_HEADERS);
|
||||||
|
assert_ne!(
|
||||||
|
header.stream_id, 2,
|
||||||
|
"second stream should have different stream_id"
|
||||||
|
);
|
||||||
|
|
||||||
|
// Simulate proxy response for second stream
|
||||||
|
let mut resp2_headers = protocol::encode_frame(
|
||||||
|
header.stream_id,
|
||||||
|
protocol::RESPONSE_HEADERS,
|
||||||
|
0,
|
||||||
|
&resp_meta,
|
||||||
|
);
|
||||||
|
hub.handle_proxy_frame(400, &mut resp2_headers).await;
|
||||||
|
|
||||||
|
let response = stream2
|
||||||
|
.wait_headers(std::time::Duration::from_secs(1))
|
||||||
|
.await
|
||||||
|
.expect("second stream should receive headers");
|
||||||
|
assert_eq!(response.status, 200);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -136,6 +136,9 @@ async fn run_proxy_reader(
|
|||||||
warn!(conn_id = conn.id, error = %e, "proxy WebSocket error");
|
warn!(conn_id = conn.id, error = %e, "proxy WebSocket error");
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
Some(Ok(Message::Ping(payload))) => {
|
||||||
|
conn.send(Message::Pong(payload));
|
||||||
|
}
|
||||||
_ => {}
|
_ => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user