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:
2026-08-29 09:36:34 +00:00
parent 5ff88756eb
commit a9ac6f6cbd
4 changed files with 223 additions and 42 deletions
+212 -19
View File
@@ -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"),
}
}
}
+1 -21
View File
@@ -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`.
+3 -2
View File
@@ -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"))]
+7
View File
@@ -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 {