mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
fix(proxy-tunnel): 修复隧道连接池竞态与跨平台兼容问题
- 修复 TCP keepalive with_retries 在 Windows 上不可用的编译问题 - dispatcher 中 try_send 失败时记录警告日志而非静默丢弃 - 优化 handler_handles 清理策略,每 64 帧定期清理 - StreamState 记住原始连接引用,清理时避免连接池竞态
This commit is contained in:
@@ -170,8 +170,9 @@ fn configure_tcp_socket(stream: &TcpStream, state: &Arc<AppState>) {
|
|||||||
if state.config.tunnel_tcp_keepalive_secs > 0 {
|
if state.config.tunnel_tcp_keepalive_secs > 0 {
|
||||||
let keepalive = socket2::TcpKeepalive::new()
|
let keepalive = socket2::TcpKeepalive::new()
|
||||||
.with_time(Duration::from_secs(state.config.tunnel_tcp_keepalive_secs))
|
.with_time(Duration::from_secs(state.config.tunnel_tcp_keepalive_secs))
|
||||||
.with_interval(Duration::from_secs(5))
|
.with_interval(Duration::from_secs(5));
|
||||||
.with_retries(3);
|
#[cfg(not(target_os = "windows"))]
|
||||||
|
let keepalive = keepalive.with_retries(3);
|
||||||
if let Err(e) = sock_ref.set_tcp_keepalive(&keepalive) {
|
if let Err(e) = sock_ref.set_tcp_keepalive(&keepalive) {
|
||||||
warn!(error = %e, "failed to set TCP keepalive on tunnel socket");
|
warn!(error = %e, "failed to set TCP keepalive on tunnel socket");
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -37,6 +37,7 @@ where
|
|||||||
// Track spawned stream handlers so we can wait for them on shutdown
|
// Track spawned stream handlers so we can wait for them on shutdown
|
||||||
let mut handler_handles: Vec<JoinHandle<()>> = Vec::new();
|
let mut handler_handles: Vec<JoinHandle<()>> = Vec::new();
|
||||||
let max_streams = state.config.tunnel_max_streams.unwrap_or(128) as usize;
|
let max_streams = state.config.tunnel_max_streams.unwrap_or(128) as usize;
|
||||||
|
let mut frames_since_cleanup: u32 = 0;
|
||||||
let stale_timeout = Duration::from_secs(state.config.tunnel_stale_timeout_secs);
|
let stale_timeout = Duration::from_secs(state.config.tunnel_stale_timeout_secs);
|
||||||
|
|
||||||
// Track last time we received any data to detect stale connections
|
// Track last time we received any data to detect stale connections
|
||||||
@@ -97,12 +98,20 @@ where
|
|||||||
Err(e) => {
|
Err(e) => {
|
||||||
warn!(stream_id = frame.stream_id, error = %e, "invalid request metadata");
|
warn!(stream_id = frame.stream_id, error = %e, "invalid request metadata");
|
||||||
// Use try_send to avoid blocking the read loop
|
// Use try_send to avoid blocking the read loop
|
||||||
let _ = frame_tx.try_send(Frame::new(
|
if frame_tx
|
||||||
frame.stream_id,
|
.try_send(Frame::new(
|
||||||
MsgType::StreamError,
|
frame.stream_id,
|
||||||
0,
|
MsgType::StreamError,
|
||||||
Bytes::from(format!("invalid request metadata: {e}")),
|
0,
|
||||||
));
|
Bytes::from(format!("invalid request metadata: {e}")),
|
||||||
|
))
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
warn!(
|
||||||
|
stream_id = frame.stream_id,
|
||||||
|
"writer channel full, StreamError dropped"
|
||||||
|
);
|
||||||
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@@ -112,12 +121,20 @@ where
|
|||||||
stream_id = frame.stream_id,
|
stream_id = frame.stream_id,
|
||||||
"max concurrent streams reached"
|
"max concurrent streams reached"
|
||||||
);
|
);
|
||||||
let _ = frame_tx.try_send(Frame::new(
|
if frame_tx
|
||||||
frame.stream_id,
|
.try_send(Frame::new(
|
||||||
MsgType::StreamError,
|
frame.stream_id,
|
||||||
0,
|
MsgType::StreamError,
|
||||||
Bytes::from("max concurrent streams reached"),
|
0,
|
||||||
));
|
Bytes::from("max concurrent streams reached"),
|
||||||
|
))
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
warn!(
|
||||||
|
stream_id = frame.stream_id,
|
||||||
|
"writer channel full, StreamError dropped"
|
||||||
|
);
|
||||||
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -163,7 +180,12 @@ where
|
|||||||
|
|
||||||
MsgType::Ping => {
|
MsgType::Ping => {
|
||||||
// Use try_send to avoid blocking the read loop when writer is congested
|
// Use try_send to avoid blocking the read loop when writer is congested
|
||||||
let _ = frame_tx.try_send(Frame::control(MsgType::Pong, frame.payload));
|
if frame_tx
|
||||||
|
.try_send(Frame::control(MsgType::Pong, frame.payload))
|
||||||
|
.is_err()
|
||||||
|
{
|
||||||
|
warn!("writer channel full, Pong dropped");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
MsgType::HeartbeatAck => {
|
MsgType::HeartbeatAck => {
|
||||||
@@ -180,9 +202,12 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Periodically clean up finished handles to avoid unbounded growth
|
// Periodically clean up finished handles to avoid unbounded growth.
|
||||||
if handler_handles.len() > max_streams {
|
// Trigger every 64 frames OR when the count exceeds max_streams.
|
||||||
|
frames_since_cleanup += 1;
|
||||||
|
if frames_since_cleanup >= 64 || handler_handles.len() > max_streams {
|
||||||
handler_handles.retain(|h| !h.is_finished());
|
handler_handles.retain(|h| !h.is_finished());
|
||||||
|
frames_since_cleanup = 0;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -63,7 +63,7 @@ class TunnelConnection:
|
|||||||
raise TunnelStreamError("frame send timeout (writer congested)")
|
raise TunnelStreamError("frame send timeout (writer congested)")
|
||||||
|
|
||||||
def create_stream(self, stream_id: int) -> _StreamState:
|
def create_stream(self, stream_id: int) -> _StreamState:
|
||||||
state = _StreamState(stream_id)
|
state = _StreamState(stream_id, conn=self)
|
||||||
self._pending_streams[stream_id] = state
|
self._pending_streams[stream_id] = state
|
||||||
return state
|
return state
|
||||||
|
|
||||||
@@ -109,9 +109,10 @@ class _StreamState:
|
|||||||
"_body_chunks",
|
"_body_chunks",
|
||||||
"_done_event",
|
"_done_event",
|
||||||
"_error",
|
"_error",
|
||||||
|
"_conn",
|
||||||
)
|
)
|
||||||
|
|
||||||
def __init__(self, stream_id: int) -> None:
|
def __init__(self, stream_id: int, conn: TunnelConnection | None = None) -> None:
|
||||||
self.stream_id = stream_id
|
self.stream_id = stream_id
|
||||||
self.status: int = 0
|
self.status: int = 0
|
||||||
self.headers: list[list[str]] = []
|
self.headers: list[list[str]] = []
|
||||||
@@ -119,6 +120,7 @@ class _StreamState:
|
|||||||
self._body_chunks: asyncio.Queue[bytes | None] = asyncio.Queue()
|
self._body_chunks: asyncio.Queue[bytes | None] = asyncio.Queue()
|
||||||
self._done_event = asyncio.Event()
|
self._done_event = asyncio.Event()
|
||||||
self._error: str | None = None
|
self._error: str | None = None
|
||||||
|
self._conn = conn
|
||||||
|
|
||||||
def set_response_headers(self, status: int, headers: list[list[str]] | dict[str, str]) -> None:
|
def set_response_headers(self, status: int, headers: list[list[str]] | dict[str, str]) -> None:
|
||||||
self.status = status
|
self.status = status
|
||||||
|
|||||||
@@ -95,7 +95,10 @@ class TunnelTransport(httpx.AsyncBaseTransport):
|
|||||||
def _cleanup_stream(self, manager: TunnelManager, stream_state: _StreamState | None) -> None:
|
def _cleanup_stream(self, manager: TunnelManager, stream_state: _StreamState | None) -> None:
|
||||||
if stream_state is None:
|
if stream_state is None:
|
||||||
return
|
return
|
||||||
conn = manager.get_connection(self._node_id)
|
# 优先从 stream 记住的原始连接上移除,避免连接池竞态
|
||||||
|
conn = stream_state._conn
|
||||||
|
if conn is None:
|
||||||
|
conn = manager.get_connection(self._node_id)
|
||||||
if conn:
|
if conn:
|
||||||
conn.remove_stream(stream_state.stream_id)
|
conn.remove_stream(stream_state.stream_id)
|
||||||
|
|
||||||
@@ -120,8 +123,10 @@ class TunnelResponseStream(httpx.AsyncByteStream):
|
|||||||
yield chunk
|
yield chunk
|
||||||
|
|
||||||
async def aclose(self) -> None:
|
async def aclose(self) -> None:
|
||||||
# 确保 stream 从 connection 的 pending 列表中移除,防止内存泄漏
|
# 从 stream 记住的原始连接上精确移除,避免连接池竞态
|
||||||
conn = self._manager.get_connection(self._node_id)
|
conn = self._stream_state._conn
|
||||||
|
if conn is None:
|
||||||
|
conn = self._manager.get_connection(self._node_id)
|
||||||
if conn:
|
if conn:
|
||||||
conn.remove_stream(self._stream_state.stream_id)
|
conn.remove_stream(self._stream_state.stream_id)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user