i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-12 20:19:51 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(ws): 用 ResponsesTurnState 收敛连接 turn 状态

Browse Source 
评审第 2 条：BoundResponsesConnection 用 response_in_flight、active_turn、
active_response_create 三个可独立变化的字段编码同一件事，8 种组合里只有 3 种
合法，非法组合只能靠调用点的 if 和「记得同时改另外两个字段」来避免。

三字段合并为一个 ResponsesTurnState：

  Idle                                  没有进行中的 logical turn
  Responding { logical, attempt }        logical 与 attempt 必须同时存在
  Replanning { logical }                 attempt 已取走去结算/重绑，logical 仍在

Replanning 不是新概念：配额透明重试期间现状就处于这个状态，只是靠
Option::take 意外得到。转换只能走 begin / detach_attempt / resume / end，
response_in_flight 与「是否接受新 response.create」都由变体推导。

由此消除的运行时不变量（原来全靠调用点自觉）：
- 有 attempt 必有 logical turn
- response_in_flight 与 attempt 同生共死（原来 client 写失败后
  active_turn=None 而 response_in_flight 仍为 true）
- logical turn 结束时必须清 attempt：原来 `active_response_create = None`
  在 connection.rs 里手写 13 处，漏一处就残留；现在只有 end() 一个出口
- 上游绑定返回的连接不再自带 response_in_flight=true 的半成品状态

同时删除 update_response_in_flight：Started 帧把已经是 true 的字段再设一次，
Close 帧因为没有解析出的 frame 而根本不触发，是纯冗余写；它在 Idle 态收到
Started 帧时还会把 response_in_flight 置真，从而永久阻塞后续 response.create。

行为等价。ActiveResponsesWebSocketRequest 改名 LogicalTurn 并随状态机移入
新的 turn_state.rs；状态机对 attempt 类型泛型化，测试用轻量替身驱动同一套
转换逻辑，无需 AppState 或真实 socket。
This commit is contained in:
AAEE86
2026-08-17 14:52:57 +08:00
committed by ZheFox
parent 9d80281b53
commit 1c5ee5228c
10 changed files with 428 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/proxy/websocket/responses/connection.rs
+22 -55
View File
@@ -21,8 +21,7 @@ use super::quota::{
send_previous_response_not_found, should_request_full_continuation_retry,
};
use super::relay_policy::{
classify_quota_relay, classify_upstream_frame, fatal_relay_policy, FatalRelaySignal,
QuotaRelayAction, QuotaRelayFacts, UpstreamFrameAction, UpstreamFrameKind,
classify_quota_relay, fatal_relay_policy, FatalRelaySignal, QuotaRelayAction, QuotaRelayFacts,
};
use super::state::BoundResponsesConnection;
use super::turn::{
@@ -65,7 +64,7 @@ pub(super) async fn relay_bound_connection(
tokio::pin!(connection_deadline);
loop {
let active_turn_deadline = bound.active_turn.as_ref().map(|turn| turn.deadline());
let active_turn_deadline = bound.turn_state.attempt().map(|turn| turn.deadline());
tokio::select! {
_ = &mut connection_deadline => {
finalize_active_turn(
@@ -79,7 +78,6 @@ pub(super) async fn relay_bound_connection(
"websocket_connection_limit_reached",
"WebSocket connection duration limit reached; reconnect to continue",
).await;
bound.active_response_create = None;
close_bound_upstream(bound).await;
close_client_socket(client_socket, CLOSE_TRY_AGAIN, "connection_limit_reached").await;
break;
@@ -105,7 +103,6 @@ pub(super) async fn relay_bound_connection(
turn_deadline.phase.error_code(),
turn_deadline.phase.client_message(),
).await;
bound.active_response_create = None;
close_bound_upstream(bound).await;
close_client_socket(
client_socket,
@@ -129,7 +126,6 @@ pub(super) async fn relay_bound_connection(
state,
ResponsesWebSocketTurnOutcome::connection_admission_lost(),
).await;
bound.active_response_create = None;
close_bound_upstream(bound).await;
send_gateway_error_with_status(
client_socket,
@@ -147,7 +143,6 @@ pub(super) async fn relay_bound_connection(
state,
ResponsesWebSocketTurnOutcome::client_disconnected(),
).await;
bound.active_response_create = None;
close_bound_upstream(bound).await;
break;
};
@@ -165,7 +160,6 @@ pub(super) async fn relay_bound_connection(
state,
ResponsesWebSocketTurnOutcome::client_disconnected(),
).await;
bound.active_response_create = None;
close_bound_upstream(bound).await;
break;
};
@@ -200,7 +194,6 @@ pub(super) async fn relay_bound_connection(
code,
"Gateway could not forward the WebSocket event upstream",
).await;
bound.active_response_create = None;
close_bound_upstream(bound).await;
close_client_socket(client_socket, CLOSE_INTERNAL_ERROR, code).await;
break;
@@ -214,7 +207,6 @@ pub(super) async fn relay_bound_connection(
state,
ResponsesWebSocketTurnOutcome::upstream_closed(),
).await;
bound.active_response_create = None;
bound.upstream = None;
close_client_socket(client_socket, 1000, "upstream_closed").await;
break;
@@ -239,7 +231,6 @@ pub(super) async fn relay_bound_connection(
"responses_websocket_receive_failed",
"Provider connection closed unexpectedly",
).await;
bound.active_response_create = None;
bound.upstream = None;
close_client_socket(client_socket, CLOSE_INTERNAL_ERROR, "upstream_receive_failed").await;
break;
@@ -268,7 +259,7 @@ pub(super) async fn relay_bound_connection(
chunked = parsed_upstream_frame
.as_ref()
.is_some_and(ParsedResponsesWebSocketFrame::is_chunked),
active_turn = bound.active_turn.is_some(),
active_turn = bound.turn_state.response_in_flight(),
"gateway received Responses WebSocket event"
);
}
@@ -315,11 +306,11 @@ pub(super) async fn relay_bound_connection(
let adapter = bound.adapter;
match parsed_upstream_frame.as_ref() {
Some(frame) => bound
.active_turn
.as_mut()
.turn_state
.attempt_mut()
.and_then(|turn| turn.observe_upstream_frame(frame, adapter)),
None => {
if let Some(turn) = bound.active_turn.as_mut() {
if let Some(turn) = bound.turn_state.attempt_mut() {
turn.observe_invalid_upstream_text(text.as_str())
}
else {
@@ -330,13 +321,12 @@ pub(super) async fn relay_bound_connection(
}
_ => None,
};
update_response_in_flight(bound, parsed_upstream_frame.as_ref());
if matches!(
observation,
Some(ResponsesWebSocketTurnObservation::Started)
| Some(ResponsesWebSocketTurnObservation::Terminal(_))
) {
if let Some(turn) = bound.active_turn.as_mut() {
if let Some(turn) = bound.turn_state.attempt_mut() {
turn.mark_stream_started(state).await;
}
}
@@ -344,9 +334,6 @@ pub(super) async fn relay_bound_connection(
Some(ResponsesWebSocketTurnObservation::Terminal(outcome)) => Some(outcome),
_ => None,
};
if terminal_outcome.is_some() {
bound.response_in_flight = false;
}
if matches!(&upstream_message, WreqWsMessage::Text(_))
&& parsed_upstream_frame.is_none()
{
@@ -359,7 +346,6 @@ pub(super) async fn relay_bound_connection(
),
)
.await;
bound.active_response_create = None;
send_responses_websocket_error(
client_socket,
policy.status_code,
@@ -380,7 +366,7 @@ pub(super) async fn relay_bound_connection(
let is_close = matches!(upstream_message, WreqWsMessage::Close(_));
let drain_for_adapter = adapter_drain_ready(
bound.pending_adapter_drain,
bound.response_in_flight,
bound.turn_state.response_in_flight(),
observation,
is_close,
);
@@ -396,7 +382,8 @@ pub(super) async fn relay_bound_connection(
};
let mut quota_relay_action = classify_quota_relay(quota_facts);
if matches!(quota_relay_action, QuotaRelayAction::AttemptTransparentRetry) {
let mut retry_turn = bound.active_turn.take().map(ActiveResponsesWebSocketTurn::disarm);
// detach_attempt 保留 logical turn：重试是同一轮请求的下一个 attempt。
let mut retry_turn = bound.turn_state.detach_attempt().map(ActiveResponsesWebSocketTurn::disarm);
if let Some(turn) = retry_turn.as_mut() {
turn.release_admission().await;
}
@@ -414,7 +401,16 @@ pub(super) async fn relay_bound_connection(
}
continue;
}
bound.active_turn = retry_turn.map(|turn| ActiveResponsesWebSocketTurn::new(state, turn));
if let Some(turn) = retry_turn {
let restored = bound
.turn_state
.resume(ActiveResponsesWebSocketTurn::new(state, turn));
debug_assert!(
restored.is_ok(),
"a failed transparent retry must be able to restore its attempt"
);
drop(restored);
}
quota_relay_action = classify_quota_relay(QuotaRelayFacts {
retry_current_turn: false,
transparent_retry_failed: true,
@@ -437,7 +433,7 @@ pub(super) async fn relay_bound_connection(
error_code = "previous_response_not_found",
"gateway will ask the client to retry the continuation with complete input"
);
let mut turn = bound.active_turn.take().map(ActiveResponsesWebSocketTurn::disarm);
let mut turn = bound.turn_state.end().map(ActiveResponsesWebSocketTurn::disarm);
if let Some(active_turn) = turn.as_mut() {
active_turn.release_admission().await;
}
@@ -453,7 +449,6 @@ pub(super) async fn relay_bound_connection(
)
.await;
}
bound.active_response_create = None;
detach_exhausted_upstream(bound, directive, &context.trace_id).await;
continue;
}
@@ -475,7 +470,6 @@ pub(super) async fn relay_bound_connection(
"Provider connection closed after reporting exhausted quota; send a new response.create to select another Provider connection",
)
.await;
bound.active_response_create = None;
detach_exhausted_upstream(bound, directive, &context.trace_id).await;
continue;
}
@@ -497,18 +491,16 @@ pub(super) async fn relay_bound_connection(
state,
ResponsesWebSocketTurnOutcome::client_disconnected(),
).await;
bound.active_response_create = None;
close_bound_upstream(bound).await;
break;
}
if let (Some(turn), Some(frame)) =
(bound.active_turn.as_mut(), parsed_upstream_frame.as_ref())
(bound.turn_state.attempt_mut(), parsed_upstream_frame.as_ref())
{
turn.capture_client_frame(frame.event());
}
if let Some(outcome) = terminal_outcome {
finalize_active_turn(bound, state, outcome).await;
bound.active_response_create = None;
} else if is_close {
finalize_active_turn(
bound,
@@ -521,7 +513,6 @@ pub(super) async fn relay_bound_connection(
let directive = bound
.pending_adapter_drain
.expect("adapter drain state should be present");
bound.active_response_create = None;
detach_exhausted_upstream(bound, directive, &context.trace_id).await;
continue;
}
@@ -549,27 +540,3 @@ async fn wait_for_connection_permit_loss(permit: Option<&aether_runtime::Admissi
}
}
fn update_response_in_flight(
bound: &mut BoundResponsesConnection,
frame: Option<&ParsedResponsesWebSocketFrame<'_>>,
) {
let Some(frame) = frame else {
return;
};
let frame_kind = if frame.is_terminal() {
UpstreamFrameKind::Terminal
} else if frame.is_started() {
UpstreamFrameKind::Started
} else {
UpstreamFrameKind::Other
};
match classify_upstream_frame(frame_kind) {
UpstreamFrameAction::Continue if frame.is_started() => {
bound.response_in_flight = true;
}
UpstreamFrameAction::FinalizeTurn => {
bound.response_in_flight = false;
}
UpstreamFrameAction::Continue | UpstreamFrameAction::FinalizeAndClose => {}
}
}