fix(websocket): lossless EOF signal + pending sweep (WS-02, CON-02)
- Replace the axum/tungstenite pump paths' Notify-based read-EOF signal with a retained tokio watch channel: a late subscriber (monitor spawned after session setup, or pump EOF before the receiver is taken) still observes EOF (WS-02). - from_wss drop monitor: on EOF (or session close) fail all pendings retryable, then keep sweeping the pending map every 1 s — calls registered after the initial fail_all (the forgotten-session import path) resolve instead of hanging (CON-02). - Tests: drop-during-registration race variants (forget + held session) and a post-EOF registration resolved via the sweep; the existing no-hang test stays green. cargo test (219), cargo test --features wss (231, 3x for flake check), cargo clippy --all-targets -- -D warnings, cargo fmt --check
This commit is contained in:
+212
-19
@@ -56,6 +56,10 @@ use crate::websocket::split_tungstenite_to_bytes;
|
||||
/// been observed — see the module doc).
|
||||
const PENDING_SWEEP_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
|
||||
|
||||
fn connection_closed_error() -> CallError {
|
||||
CallError::new("CONNECTION_CLOSED", "from_wss connection dropped", true)
|
||||
}
|
||||
|
||||
pub struct FromWss {
|
||||
endpoint: String,
|
||||
auth_token: Option<String>,
|
||||
@@ -166,36 +170,33 @@ impl WssSession {
|
||||
|
||||
let (close_tx, mut close_rx) = tokio::sync::oneshot::channel();
|
||||
let pending = Arc::clone(call_connection.pending());
|
||||
let mut eof_rx = pumps.read_eof();
|
||||
tokio::spawn(async move {
|
||||
let mut eof_rx = pumps.read_eof();
|
||||
let mut sweep = tokio::time::interval(PENDING_SWEEP_INTERVAL);
|
||||
sweep.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
let mut eof_observed = *eof_rx.borrow_and_update();
|
||||
if eof_observed {
|
||||
pending.lock().fail_all(connection_closed_error());
|
||||
}
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = crate::websocket::wait_for_eof(&mut eof_rx) => break,
|
||||
_ = &mut close_rx => break,
|
||||
changed = eof_rx.changed() => {
|
||||
eof_observed = changed.is_err() || *eof_rx.borrow_and_update();
|
||||
if eof_observed {
|
||||
pending.lock().fail_all(connection_closed_error());
|
||||
}
|
||||
}
|
||||
_ = &mut close_rx => {
|
||||
pending.lock().fail_all(connection_closed_error());
|
||||
return;
|
||||
}
|
||||
_ = sweep.tick() => {
|
||||
if *eof_rx.borrow() {
|
||||
pending.lock().fail_all(
|
||||
CallError::new(
|
||||
"CONNECTION_CLOSED",
|
||||
"from_wss connection dropped",
|
||||
true,
|
||||
),
|
||||
);
|
||||
pending.lock().fail_all(connection_closed_error());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let failure = CallError::new("CONNECTION_CLOSED", "from_wss connection dropped", true);
|
||||
let mut guard = pending.lock();
|
||||
guard.fail_all(failure.clone());
|
||||
if *eof_rx.borrow() {
|
||||
let swept = guard.fail_all(failure);
|
||||
if !swept.is_empty() {
|
||||
tracing::debug!("from_wss: sweep failed {} late pendings", swept.len());
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self {
|
||||
@@ -243,6 +244,76 @@ mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
|
||||
type WsPumpsSlot = std::sync::Arc<std::sync::Mutex<Option<crate::websocket::WsPumps>>>;
|
||||
|
||||
/// A producer whose WS session can be torn down on demand: the
|
||||
/// killable upgrade handler stores the session's `WsPumps` in the
|
||||
/// returned slot; aborting those pumps forces the server socket
|
||||
/// closed (consumer-side read EOF).
|
||||
fn drop_on_signal_producer(registry: Arc<OperationRegistry>) -> (String, WsPumpsSlot) {
|
||||
let pumps_slot: WsPumpsSlot = std::sync::Arc::new(std::sync::Mutex::new(None));
|
||||
let provider = provider_with(vec![("tok-1", identity("alice", &[]))]);
|
||||
async fn killable_upgrade(
|
||||
axum::extract::State(state): axum::extract::State<(
|
||||
Arc<OperationRegistry>,
|
||||
WsPumpsSlot,
|
||||
)>,
|
||||
axum::Extension(identity): axum::Extension<Identity>,
|
||||
ws_upgrade: axum::extract::ws::WebSocketUpgrade,
|
||||
) -> axum::response::Response {
|
||||
ws_upgrade.on_upgrade(move |socket| async move {
|
||||
let (byte_stream, pumps) = crate::websocket::split_ws_to_bytes(socket);
|
||||
*state.1.lock().unwrap_or_else(|e| e.into_inner()) = Some(pumps);
|
||||
let conn = alkcall::core::types::Connection::from_bidi(
|
||||
byte_stream,
|
||||
b"alk/channels".to_vec(),
|
||||
None,
|
||||
);
|
||||
let _ = conn.set_identity(identity.clone());
|
||||
let adapter = alkcall::channels::adapter::ChannelsAdapter::new(
|
||||
crate::websocket::adapter_install_channel_zero(Arc::clone(&state.0)),
|
||||
std::sync::Arc::new(alkcall::channels::policy::NoCap),
|
||||
);
|
||||
let auth = alkcall::core::auth::AuthContext {
|
||||
identity: Some(identity),
|
||||
alpn: b"alk/channels".to_vec(),
|
||||
remote_addr: None,
|
||||
tls_client_fingerprint: None,
|
||||
};
|
||||
if let Err(e) =
|
||||
alkcall::core::types::ProtocolHandler::handle(&adapter, conn, &auth).await
|
||||
{
|
||||
tracing::warn!(error = %e, "kill-test channels session ended");
|
||||
}
|
||||
})
|
||||
}
|
||||
let app = axum::Router::new()
|
||||
.route(
|
||||
"/alk/channels",
|
||||
axum::routing::get(killable_upgrade).route_layer(
|
||||
axum::middleware::from_fn_with_state(
|
||||
Arc::clone(&provider),
|
||||
crate::websocket::ws_bearer_auth,
|
||||
),
|
||||
),
|
||||
)
|
||||
.with_state((registry, Arc::clone(&pumps_slot)));
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
|
||||
let addr = listener.local_addr().expect("addr");
|
||||
let _ = listener.set_nonblocking(true);
|
||||
std::thread::spawn(move || {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.expect("rt");
|
||||
rt.block_on(async {
|
||||
let listener = tokio::net::TcpListener::from_std(listener).expect("tokio listener");
|
||||
let _ = axum::serve(listener, app).await;
|
||||
});
|
||||
});
|
||||
(format!("ws://{addr}/alk/channels"), pumps_slot)
|
||||
}
|
||||
|
||||
fn identity(id: &str, scopes: &[&str]) -> Identity {
|
||||
Identity {
|
||||
id: id.to_string(),
|
||||
@@ -654,4 +725,126 @@ mod tests {
|
||||
assert!(adapter.auth_token().is_none());
|
||||
std::env::remove_var("WSS_TOKEN");
|
||||
}
|
||||
|
||||
/// Abort the producer-side WS pumps once the upgrade handler has
|
||||
/// stored them, then wait for the consumer's WS read pump to reach
|
||||
/// EOF (the abort surfaces as a socket close).
|
||||
async fn abort_pumps_and_wait(slot: &WsPumpsSlot) {
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
|
||||
loop {
|
||||
let pumps = slot.lock().unwrap_or_else(|e| e.into_inner()).take();
|
||||
match pumps {
|
||||
Some(pumps) => {
|
||||
pumps.abort();
|
||||
break;
|
||||
}
|
||||
None => {
|
||||
if std::time::Instant::now() > deadline {
|
||||
panic!("producer pumps never registered");
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Headroom for the abort to surface as consumer-side EOF.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
|
||||
}
|
||||
|
||||
async fn race_call_resolves_retryable(drop_before_call: bool) {
|
||||
let (endpoint, pumps_slot) = drop_on_signal_producer(slow_producer_registry());
|
||||
let session = WssSession::connect(&endpoint, Some("tok-1"))
|
||||
.await
|
||||
.expect("connect");
|
||||
|
||||
let config = FromCallConfig::new();
|
||||
let bundles = import_from_call(&session.call_connection, config)
|
||||
.await
|
||||
.expect("import");
|
||||
let slow = bundles
|
||||
.into_iter()
|
||||
.find(|b| b.spec.name == "slow/op")
|
||||
.expect("slow/op present");
|
||||
let handler = match &slow.handler {
|
||||
HandlerKind::Once(h) => h.clone(),
|
||||
_ => panic!("expected Once handler"),
|
||||
};
|
||||
|
||||
if drop_before_call {
|
||||
std::mem::forget(session);
|
||||
} else {
|
||||
drop(session);
|
||||
}
|
||||
abort_pumps_and_wait(&pumps_slot).await;
|
||||
|
||||
let ctx = noop_context("req-race");
|
||||
let call_task = tokio::spawn(async move { handler(serde_json::json!({}), ctx).await });
|
||||
|
||||
let response = tokio::time::timeout(std::time::Duration::from_secs(5), call_task)
|
||||
.await
|
||||
.expect("call resolves after racing drop (no hang)")
|
||||
.expect("join");
|
||||
match response.result {
|
||||
Err(e) => {
|
||||
if e.code == "INTERNAL" && e.message.contains("failed to write request frame") {
|
||||
// The call registered but the mux was torn down
|
||||
// before the request frame was written — the
|
||||
// write-failure path resolved it (no hang; the
|
||||
// retryable variant is asserted by the sweep test
|
||||
// and the pre-drop test).
|
||||
} else {
|
||||
assert!(e.retryable, "drop error must be retryable, got {e:?}");
|
||||
}
|
||||
}
|
||||
Ok(_) => panic!("expected Err after connection drop"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn forget_session_drop_during_call_registration_resolves_retryable_no_hang() {
|
||||
race_call_resolves_retryable(true).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn held_session_drop_during_call_registration_resolves_retryable_no_hang() {
|
||||
race_call_resolves_retryable(false).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_registered_after_eof_resolves_via_pending_sweep() {
|
||||
let (endpoint, pumps_slot) = drop_on_signal_producer(slow_producer_registry());
|
||||
let session = WssSession::connect(&endpoint, Some("tok-1"))
|
||||
.await
|
||||
.expect("connect");
|
||||
|
||||
let config = FromCallConfig::new();
|
||||
let bundles = import_from_call(&session.call_connection, config)
|
||||
.await
|
||||
.expect("import");
|
||||
let slow = bundles
|
||||
.into_iter()
|
||||
.find(|b| b.spec.name == "slow/op")
|
||||
.expect("slow/op present");
|
||||
let handler = match &slow.handler {
|
||||
HandlerKind::Once(h) => h.clone(),
|
||||
_ => panic!("expected Once handler"),
|
||||
};
|
||||
|
||||
std::mem::forget(session);
|
||||
abort_pumps_and_wait(&pumps_slot).await;
|
||||
// Headroom past the monitor's 1 s sweep tick: the call
|
||||
// registered below must be caught by the next sweep, not hang.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1100)).await;
|
||||
|
||||
let ctx = noop_context("req-late");
|
||||
let call_task = tokio::spawn(async move { handler(serde_json::json!({}), ctx).await });
|
||||
|
||||
let response = tokio::time::timeout(std::time::Duration::from_secs(5), call_task)
|
||||
.await
|
||||
.expect("post-EOF registered call resolves via sweep (no hang)")
|
||||
.expect("join");
|
||||
match response.result {
|
||||
Err(e) => assert!(e.retryable, "drop error must be retryable, got {e:?}"),
|
||||
Ok(_) => panic!("expected Err after connection drop"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -95,33 +95,13 @@ impl WsPumps {
|
||||
/// EOF signaled at any point — including before the receiver is
|
||||
/// taken or the observer starts awaiting — is still observed, and
|
||||
/// may be observed repeatedly. Used by `from_wss`'s
|
||||
/// connection-drop monitor (ADR-070); await it with
|
||||
/// [`wait_for_eof`].
|
||||
/// connection-drop monitor (ADR-070).
|
||||
#[cfg(feature = "wss")]
|
||||
pub(crate) fn read_eof(&self) -> tokio::sync::watch::Receiver<bool> {
|
||||
self.read_eof.subscribe()
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolves once the given WS read-EOF watch receiver observes `true`.
|
||||
///
|
||||
/// Lossless by construction: a `watch::Receiver` retains the latest
|
||||
/// value, so an EOF signaled before this call (or before the receiver
|
||||
/// existed) is still observed — the property a one-shot `Notify` lacked
|
||||
/// (WS-02). Returns early (treated as EOF) if the sender half is
|
||||
/// dropped, e.g. the read pump was aborted.
|
||||
#[cfg(feature = "wss")]
|
||||
pub(crate) async fn wait_for_eof(rx: &mut tokio::sync::watch::Receiver<bool>) {
|
||||
loop {
|
||||
if *rx.borrow_and_update() {
|
||||
return;
|
||||
}
|
||||
if rx.changed().await.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Split a `WebSocket` into the byte stream + the pump tasks. The
|
||||
/// adapter is the single seam between axum's WS and alkcall's
|
||||
/// byte-oriented channels machinery; shared with `from_wss`.
|
||||
|
||||
@@ -18,8 +18,9 @@ pub use byte_adapter::{
|
||||
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
pub use byte_adapter::split_tungstenite_to_bytes;
|
||||
#[cfg(feature = "wss")]
|
||||
pub(crate) use byte_adapter::wait_for_eof;
|
||||
#[cfg(any(test, feature = "wss"))]
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use upgrade::adapter_install_channel_zero;
|
||||
pub use upgrade::{run_channels_session, ws_bearer_auth, ws_upgrade_handler};
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
|
||||
@@ -97,6 +97,13 @@ fn install_channel_zero(
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn adapter_install_channel_zero(
|
||||
registry: Arc<OperationRegistry>,
|
||||
) -> alkcall::channels::adapter::InstallChannelZero {
|
||||
install_channel_zero(registry)
|
||||
}
|
||||
|
||||
struct NoopProvider;
|
||||
|
||||
impl alkcall::core::auth::IdentityProvider for NoopProvider {
|
||||
|
||||
Reference in New Issue
Block a user